From 1dc10b8587b21f1af477c281dc8df850e3b57e9a Mon Sep 17 00:00:00 2001 From: Felipe Date: Sat, 1 Aug 2026 18:45:33 -0400 Subject: [PATCH 01/28] Add unit tests for atomic units and improve scratch directory handling - Introduced tests for atomic units including registration, execution, and validation. - Added a new scratch.py module to manage temporary directories for dataset caching during tests. - Updated existing tests to utilize the new scratch directory management to avoid cluttering the repository. - Ensured that unit tests cover various scenarios including validation failures and context management. --- DashAI/back/initial_components.py | 15 + DashAI/back/job/dataset_job.py | 6 + DashAI/back/job/model_job.py | 333 +++------------- DashAI/back/units/__init__.py | 3 + DashAI/back/units/base_unit.py | 106 +++++ DashAI/back/units/build_model_unit.py | 244 ++++++++++++ DashAI/back/units/context.py | 188 +++++++++ DashAI/back/units/evaluate_model_unit.py | 92 +++++ DashAI/back/units/fit_model_unit.py | 230 +++++++++++ DashAI/back/units/load_dataset_unit.py | 74 ++++ DashAI/back/units/prepare_and_split_unit.py | 182 +++++++++ DashAI/back/units/save_model_unit.py | 41 ++ .../back/api/test_model_job_orchestration.py | 372 ++++++++++++++++++ tests/back/api/test_units_api.py | 79 ++++ tests/back/conftest.py | 9 + .../base_tabular_dataloader_tests.py | 9 +- tests/back/dataloaders/test_dashai_dataset.py | 3 +- tests/back/explainers/test_explainers.py | 3 +- tests/back/explainers/test_lib_explainers.py | 3 +- tests/back/explainers/test_new_explainers.py | 3 +- tests/back/explainers/test_task_explainers.py | 3 +- .../back/models/test_bow_text_class_model.py | 3 +- .../models/test_deberta_v3_transformer.py | 3 +- .../models/test_distilbert_transformer.py | 3 +- .../models/test_modernbert_transformer.py | 3 +- .../back/models/test_tabular_class_models.py | 3 +- tests/back/scratch.py | 50 +++ tests/back/tasks/test_tasks.py | 12 +- tests/back/units/__init__.py | 0 tests/back/units/test_base_unit.py | 132 +++++++ tests/back/units/test_build_model_unit.py | 74 ++++ tests/back/units/test_context.py | 180 +++++++++ tests/back/units/test_evaluate_model_unit.py | 24 ++ tests/back/units/test_fit_model_unit.py | 40 ++ 34 files changed, 2229 insertions(+), 296 deletions(-) create mode 100644 DashAI/back/units/__init__.py create mode 100644 DashAI/back/units/base_unit.py create mode 100644 DashAI/back/units/build_model_unit.py create mode 100644 DashAI/back/units/context.py create mode 100644 DashAI/back/units/evaluate_model_unit.py create mode 100644 DashAI/back/units/fit_model_unit.py create mode 100644 DashAI/back/units/load_dataset_unit.py create mode 100644 DashAI/back/units/prepare_and_split_unit.py create mode 100644 DashAI/back/units/save_model_unit.py create mode 100644 tests/back/api/test_model_job_orchestration.py create mode 100644 tests/back/api/test_units_api.py create mode 100644 tests/back/scratch.py create mode 100644 tests/back/units/__init__.py create mode 100644 tests/back/units/test_base_unit.py create mode 100644 tests/back/units/test_build_model_unit.py create mode 100644 tests/back/units/test_context.py create mode 100644 tests/back/units/test_evaluate_model_unit.py create mode 100644 tests/back/units/test_fit_model_unit.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index e2519b6b0..1882857ff 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -346,6 +346,14 @@ from DashAI.back.tasks.text_to_text_generation_task import TextToTextGenerationTask from DashAI.back.tasks.translation_task import TranslationTask +# Units +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.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit + logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -508,6 +516,13 @@ def get_initial_components(): DatasetJob, GenerativeJob, PipelineJob, + # Units + LoadDatasetUnit, + PrepareAndSplitUnit, + BuildModelUnit, + FitModelUnit, + EvaluateModelUnit, + SaveModelUnit, # Explainers ContrastiveShap, DiceCounterfactual, diff --git a/DashAI/back/job/dataset_job.py b/DashAI/back/job/dataset_job.py index 4193dc3a3..d7caef39a 100644 --- a/DashAI/back/job/dataset_job.py +++ b/DashAI/back/job/dataset_job.py @@ -94,6 +94,7 @@ def run( import json import os import shutil + import tempfile import uuid from pathlib import Path @@ -114,6 +115,11 @@ def run( n_sample = self.kwargs.get("n_sample", None) file_path = self.kwargs.get("file_path") temp_dir = self.kwargs.get("temp_dir") + if not temp_dir: + # The dataloaders forward this path to HuggingFace as ``cache_dir``. + # Passing it along unset would stringify to "None" and create a + # directory literally named "None" in the working directory. + temp_dir = tempfile.mkdtemp(prefix="dashai-dataset-") url = self.kwargs.get("url", "") try: diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 3ffb0a799..23e2c8bcb 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -1,25 +1,23 @@ import logging -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc from sqlalchemy.orm.attributes import flag_modified -from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum -from DashAI.back.dependencies.database.models import Dataset, Metric, ModelSession, Run -from DashAI.back.dependencies.downloads.nested import missing_downloads +from DashAI.back.dependencies.database.models import ModelSession, Run from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.metrics.base_metric import BaseMetric -from DashAI.back.models.base_model import BaseModel -from DashAI.back.models.model_factory import ModelFactory -from DashAI.back.optimizers.base_optimizer import BaseOptimizer -from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker - from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset - logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -93,192 +91,60 @@ def run( ) -> None: import gc import json - import os - import pickle from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - prepare_for_model_session, - select_columns, - split_dataset, - ) - - component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] # Get the necessary parameters run_id: int = self.kwargs["run_id"] + ctx = ExecutionContext(refs={"run_id": 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.") run.huey_id = self.kwargs.get("huey_id", None) db.commit() self.report_progress(0.05, "Preparing data") try: - # Get the model session, dataset, task, metrics and splits + # The model session holds the configuration every unit reads. model_session: ModelSession = db.get(ModelSession, run.model_session_id) if not model_session: raise JobError( f"Model session {run.model_session_id} does not exist in DB." ) - dataset: Dataset = db.get(Dataset, model_session.dataset_id) - if not dataset: - raise JobError( - f"Dataset {model_session.dataset_id} does not exist in DB." - ) - - try: - loaded_dataset: "DashAIDataset" = 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: - # Get metrics from model session - train_metrics: List[BaseMetric] = [ - component_registry[m]["class"] - for m in model_session.train_metrics - ] - validation_metrics: List[BaseMetric] = [ - component_registry[m]["class"] - for m in model_session.validation_metrics - ] - test_metrics: List[BaseMetric] = [ - component_registry[m]["class"] - for m in model_session.test_metrics - ] - - except Exception as e: - log.exception(e) - raise JobError( - "Unable to find metrics associated with" - f"Task {model_session.task_name} in registry", - ) from e - - try: - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=model_session.input_columns, - output_columns=model_session.output_columns, - ) - n_labels = task.num_labels( - prepared_dataset, model_session.output_columns[0] - ) - - splits = json.loads(model_session.splits) - prepared_dataset, splits = prepare_for_model_session( - dataset=prepared_dataset, - splits=splits, - output_columns=model_session.output_columns, - ) - - run.split_indexes = json.dumps( - { - "train_indexes": splits["train_indexes"], - "test_indexes": splits["test_indexes"], - "val_indexes": splits["val_indexes"], - } - ) - - x, y = select_columns( - prepared_dataset, - model_session.input_columns, - model_session.output_columns, - ) - - x = split_dataset(x) - y = split_dataset(y) - - except Exception as e: - log.exception(e) - raise JobError( - f"""Can not prepare Dataset {dataset.id} - for Task {model_session.task_name}""", - ) from e - - 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 - if getattr(run_model_class, "REQUIRES_DOWNLOAD", False) and not ( - run_model_class.is_downloaded() - ): - raise JobError( - f"Model {run.model_name} is not downloaded. " - "Download it before training." - ) - nested_missing = missing_downloads(run.parameters, component_registry) - if nested_missing: - names = ", ".join(m["name"] for m in nested_missing) - raise JobError( - "These components are not downloaded. " - f"Download them before training: {names}." - ) - try: - factory = ModelFactory( - run_model_class, - run.parameters, - run_id, - x, - y, - train_metrics, - validation_metrics, - test_metrics, - n_labels=n_labels, - ) - model: BaseModel = factory.model - run_optimizable_parameters = factory.optimizable_parameters + LoadDatasetUnit(dataset_id=model_session.dataset_id)(ctx) + + PrepareAndSplitUnit( + task_name=model_session.task_name, + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + splits=json.loads(model_session.splits), + )(ctx) + + run.split_indexes = json.dumps(ctx.require("split_indexes")) + + # __call__ runs validate() (the download gate) before execute() + # for every unit, so no separate pre-check is needed here. + BuildModelUnit( + model={"component": run.model_name, "params": run.parameters}, + train_metrics=model_session.train_metrics, + validation_metrics=model_session.validation_metrics, + test_metrics=model_session.test_metrics, + )(ctx) + + # Resolving the optimizer before the status changes keeps an + # invalid configuration from ever reporting that training began. + fit_model = FitModelUnit( + optimizer={ + "component": run.optimizer_name, + "params": run.optimizer_parameters, + }, + goal_metric=run.goal_metric, + ) + fit_model.validate(ctx) - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to instantiate model using run {run_id}", - ) from e - try: - if run_optimizable_parameters: - goal_metric = component_registry[run.goal_metric] - except Exception as e: - log.exception(e) - raise JobError( - f"Metric is not compatible with the Task. {e}", - ) from e - try: - # Optimizer configuration - if run_optimizable_parameters: - run_optimizer_class = component_registry[run.optimizer_name][ - "class" - ] - optimizer: BaseOptimizer = run_optimizer_class( - **run.optimizer_parameters - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error instantiating optimizer {run.optimizer_name}, {e}", - ) from e try: run.set_status_as_started() db.commit() @@ -288,57 +154,16 @@ def run( "Connection with the database failed", ) from e self.report_progress(0.2, "Training") - try: - # Hyperparameter Tunning - plot_paths = [] - if not run_optimizable_parameters: - model.train( - x["train"], y["train"], x["validation"], y["validation"] - ) - else: - optimizer.optimize( - model, - x, - y, - run_optimizable_parameters, - goal_metric, - task, - ) - model = optimizer.get_model() - best_params = optimizer.get_best_params() - - old_parameters = run.parameters.copy() - updated_parameters = factory.update_parameters( - old_parameters, best_params - ) - - run.parameters = updated_parameters - flag_modified(run, "parameters") - db.commit() - - # Generate hyperparameter plot - from DashAI.back.core.artifacts import normalize_artifacts - - trials = optimizer.get_trials_values() - plot_filenames, plots = optimizer.create_plots( - trials, - run_id, - n_params=len(run_optimizable_parameters), - goal_metric=goal_metric, - ) - normalized_plots = normalize_artifacts(plots) - for filename, plot in zip( - plot_filenames, normalized_plots, strict=False - ): - plot_path = os.path.join(config["RUNS_PATH"], filename) - with open(plot_path, "wb") as file: - pickle.dump(plot, file) - plot_paths.append(plot_path) - except Exception as e: - log.exception(e) - raise JobError( - f"Model training failed {e}", - ) from e + + fit_model(ctx) + + plot_paths = ctx.require("plot_paths") + + if ctx.has("best_parameters"): + run.parameters = ctx.get("best_parameters") + flag_modified(run, "parameters") + db.commit() + try: paths = plot_paths + [None] * (4 - len(plot_paths)) ( @@ -355,61 +180,16 @@ def run( ) from e self.report_progress(0.85, "Computing metrics") - # Calculate metrics at the end of training if not done already - try: - last_train_metric = ( - db.query(Metric) - .filter_by(run_id=run.id, split="TRAIN", level="LAST") - .first() - ) - if not last_train_metric: - model.calculate_metrics( - split=SplitEnum.TRAIN, - level=LevelEnum.LAST, - ) - last_val_metric = ( - db.query(Metric) - .filter_by(run_id=run.id, split="VALIDATION", level="LAST") - .first() - ) - if not last_val_metric: - model.calculate_metrics( - split=SplitEnum.VALIDATION, - level=LevelEnum.LAST, - ) - last_test_metric = ( - db.query(Metric) - .filter_by(run_id=run.id, split="TEST", level="LAST") - .first() - ) - if not last_test_metric: - model.calculate_metrics( - split=SplitEnum.TEST, - level=LevelEnum.LAST, - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Metric calculation failed {e}", - ) from e + EvaluateModelUnit()(ctx) self.report_progress(0.95, "Saving model") - try: - run_path = os.path.join(config["RUNS_PATH"], str(run.id)) - model.save(run_path) - except Exception as e: - log.exception(e) - raise JobError( - "Model saving failed", - ) from e + SaveModelUnit()(ctx) try: - run.run_path = run_path + run.run_path = ctx.require("model_path") db.commit() except exc.SQLAlchemyError as e: log.exception(e) - run.set_status_as_error() - db.commit() raise JobError( "Connection with the database failed", ) from e @@ -426,4 +206,5 @@ def run( db.commit() raise e finally: + ctx.clear_cache() gc.collect() diff --git a/DashAI/back/units/__init__.py b/DashAI/back/units/__init__.py new file mode 100644 index 000000000..f9cc3fccc --- /dev/null +++ b/DashAI/back/units/__init__.py @@ -0,0 +1,3 @@ +# flake8: noqa +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError diff --git a/DashAI/back/units/base_unit.py b/DashAI/back/units/base_unit.py new file mode 100644 index 000000000..3bb5c2947 --- /dev/null +++ b/DashAI/back/units/base_unit.py @@ -0,0 +1,106 @@ +"""Base class for the atomic units a job is composed of.""" + +import logging +from abc import ABCMeta, abstractmethod +from typing import Final, Tuple, final + +from DashAI.back.config_object import ConfigObject +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.context import ExecutionContext + +logger = logging.getLogger(__name__) + + +class BaseUnit(ConfigObject, metaclass=ABCMeta): + """Abstract class for all atomic units. + + A unit is the smallest reusable piece of a job: it declares the context + keys it needs, the keys it produces, and does one thing. Jobs compose units + into a sequence; the orchestration around them (database transactions, + status transitions, progress reporting) stays in the job. + + Units never read nor mutate the ``Run`` row. ``run_id`` travels through the + context as an opaque correlation id because ``ModelFactory`` needs it for + ``BaseModel.calculate_metrics`` to work, but the ownership of the row + belongs to the job. + + ``BaseUnit`` deliberately does not inherit from ``BaseJob``: the registry + derives a component's type by walking the MRO for ancestors whose name + contains "Base" and that declare a ``TYPE``, and it rejects components with + more than one candidate. Inheriting from both would make registration fail. + """ + + TYPE: Final[str] = "Unit" + + #: Context keys that must be present before the unit runs. + REQUIRES: Tuple[str, ...] = () + #: Context keys the unit guarantees after it runs. + PROVIDES: Tuple[str, ...] = () + + SCHEMA: BaseSchema = BaseSchema + + def __init__(self, **config) -> None: + """Store the unit configuration. + + Parameters + ---------- + config : dict + Configuration of the unit, as declared by its schema. + """ + self.config = config + + def validate(self, ctx: ExecutionContext) -> None: + """Check preconditions without executing the unit. + + Runs before the job commits to any observable state change, so a unit + can reject an impossible configuration early. No-op by default. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + """ + + @abstractmethod + def execute(self, ctx: ExecutionContext) -> None: + """Do the unit's work, reading from and writing to the context. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + """ + raise NotImplementedError + + @final + def __call__(self, ctx: ExecutionContext) -> None: + """Run the unit, enforcing its declared contract. + + Calls ``validate`` before ``execute`` so a caller that just does + ``unit(ctx)`` — the sanctioned way to run a unit — always gets its + precondition checks (e.g. a download gate) for free. An orchestrator + that needs ``validate`` to run earlier, ahead of some other state + change, is still free to call ``unit.validate(ctx)`` directly first; + ``validate`` runs again here, which is redundant but harmless. + + Parameters + ---------- + ctx : ExecutionContext + The shared execution context. + + Raises + ------ + UnitContractError + If a required key is missing before execution or a promised key is + missing after it. + """ + for key in self.REQUIRES: + ctx.require(key) + + self.validate(ctx) + + logger.debug("Running unit %s", type(self).__name__) + self.execute(ctx) + + for key in self.PROVIDES: + ctx.require(key) diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py new file mode 100644 index 000000000..45c14c936 --- /dev/null +++ b/DashAI/back/units/build_model_unit.py @@ -0,0 +1,244 @@ +"""Unit that instantiates a model with its parameters, data and metrics.""" + +import logging +from typing import TYPE_CHECKING, List + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.downloads.nested import missing_downloads +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.metrics.base_metric import BaseMetric + from DashAI.back.models.base_model import BaseModel + +log = logging.getLogger(__name__) + + +def _metrics_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field()), + placeholder=[], + description=description, + alias=alias, + ) + + +class BuildModelSchema(BaseSchema): + model: schema_field( + component_field(parent="BaseModel"), + placeholder={"component": "SVC", "params": {}}, + description=MultilingualString( + en="Model to instantiate, along with its own configuration.", + es="Modelo a instanciar, junto con su propia configuración.", + pt="Modelo a instanciar, junto com a sua própria configuração.", + de="Zu instanziierendes Modell samt seiner eigenen Konfiguration.", + zh="要实例化的模型及其自身配置。", + ), + alias=MultilingualString( + en="Model", es="Modelo", pt="Modelo", de="Modell", zh="模型" + ), + ) # type: ignore + train_metrics: _metrics_field( + alias=MultilingualString( + en="Train metrics", + es="Métricas de entrenamiento", + pt="Métricas de treino", + de="Trainingsmetriken", + zh="训练指标", + ), + description=MultilingualString( + en="Metrics evaluated on the train split.", + es="Métricas evaluadas sobre la partición de entrenamiento.", + pt="Métricas avaliadas na partição de treino.", + de="Auf der Trainingsteilmenge ausgewertete Metriken.", + zh="在训练集上评估的指标。", + ), + ) # type: ignore + validation_metrics: _metrics_field( + alias=MultilingualString( + en="Validation metrics", + es="Métricas de validación", + pt="Métricas de validação", + de="Validierungsmetriken", + zh="验证指标", + ), + description=MultilingualString( + en="Metrics evaluated on the validation split.", + es="Métricas evaluadas sobre la partición de validación.", + pt="Métricas avaliadas na partição de validação.", + de="Auf der Validierungsteilmenge ausgewertete Metriken.", + zh="在验证集上评估的指标。", + ), + ) # type: ignore + test_metrics: _metrics_field( + alias=MultilingualString( + en="Test metrics", + es="Métricas de prueba", + pt="Métricas de teste", + de="Testmetriken", + zh="测试指标", + ), + description=MultilingualString( + en="Metrics evaluated on the test split.", + es="Métricas evaluadas sobre la partición de prueba.", + pt="Métricas avaliadas na partição de teste.", + de="Auf der Testteilmenge ausgewertete Metriken.", + zh="在测试集上评估的指标。", + ), + ) # type: ignore + + +class BuildModelUnit(BaseUnit): + """Instantiate an untrained model bound to its data and metrics. + + ``ModelFactory`` attaches the run id, the data splits and the metric + classes to the model instance, which is what later lets the model log + metrics on its own during and after training. The metrics are configured + here rather than in the evaluation unit because models use them *while* + training to log at the step and epoch levels. + + ``validate`` checks that the model and every component nested in its + parameters have been downloaded, so an impossible run is rejected before + anything observable happens. + + The model is configured as a component field, the same way a model picks + its own sub-components: the value is ``{"component": , "params": + {...}}``. The parameter tree stays an opaque ``dict`` in the schema and the + front resolves it recursively, fetching the chosen component's schema to + render the nested form. ``ModelFactory`` walks the same shape to build the + object graph. + """ + + SCHEMA = BuildModelSchema + + REQUIRES = ("x", "y", "n_labels") + PROVIDES = ("model", "factory", "optimizable_parameters", "model_parameters") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + @property + def model_name(self) -> str: + return self.config["model"]["component"] + + @property + def model_parameters(self) -> dict: + return self.config["model"]["params"] + + def _resolve_model_class(self) -> type: + """Resolve the model class from the registry, memoized on this unit. + + Memoized on the unit instance, not in the shared context: a context + can outlive a single ``BuildModelUnit`` (a future DAG could have more + than one build-model node feeding the same run), and a context-global + cache key would make the second instance silently reuse the first + one's model class. + """ + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + model_name: str = self.model_name + + 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 validate(self, ctx: ExecutionContext) -> None: + from kink import di + + component_registry = di["component_registry"] + model_name: str = self.model_name + parameters = self.model_parameters + + model_class = self._resolve_model_class() + + if getattr(model_class, "REQUIRES_DOWNLOAD", False) and not ( + model_class.is_downloaded() + ): + raise JobError( + f"Model {model_name} is not downloaded. Download it before training." + ) + + nested_missing = missing_downloads(parameters, component_registry) + if nested_missing: + names = ", ".join(m["name"] for m in nested_missing) + raise JobError( + "These components are not downloaded. " + f"Download them before training: {names}." + ) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.models.model_factory import ModelFactory + + component_registry = di["component_registry"] + + parameters = self.model_parameters + + model_class = self._resolve_model_class() + + try: + train_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] for m in self.config["train_metrics"] + ] + validation_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] + for m in self.config["validation_metrics"] + ] + test_metrics: List["BaseMetric"] = [ + component_registry[m]["class"] for m in self.config["test_metrics"] + ] + except Exception as e: + log.exception(e) + raise JobError( + "Unable to find metrics associated with" + f"Task {ctx.get('task_name')} in registry", + ) from e + + try: + factory = ModelFactory( + model_class, + parameters, + ctx.get("run_id"), + ctx.require("x"), + ctx.require("y"), + train_metrics, + validation_metrics, + test_metrics, + n_labels=ctx.require("n_labels"), + ) + model: "BaseModel" = factory.model + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to instantiate model using run {ctx.get('run_id')}", + ) from e + + # The original tree is what the search unit rewrites with the best + # values found, so it travels as a reference instead of being read + # from this unit's configuration again. + ctx.put_ref("model_parameters", parameters) + ctx.put("factory", factory) + ctx.put("model", model) + ctx.put("optimizable_parameters", factory.optimizable_parameters) diff --git a/DashAI/back/units/context.py b/DashAI/back/units/context.py new file mode 100644 index 000000000..7782ec93f --- /dev/null +++ b/DashAI/back/units/context.py @@ -0,0 +1,188 @@ +"""Execution context shared between atomic units.""" + +import copy +import json +from typing import Any, Dict, Optional + + +class UnitContractError(Exception): + """Raised when a unit's declared contract is violated. + + Either a required key is missing from the context before execution, or a + promised key is missing after it. + """ + + +class ExecutionContext: + """State that crosses the boundary between atomic units. + + The context deliberately splits its state in two halves with different + rules: + + * ``refs``: JSON serializable references (ids, paths, index lists). This is + the only half that can cross a process boundary. Jobs are shipped to the + Huey worker with dill and rebuild their dependencies from a fresh + container, so anything that must survive that trip has to be expressible + as plain data. + * ``cache``: live objects (datasets, models, tasks). Never serialized; + always derivable again from the refs. + + Keeping the halves apart is what allows the same unit to run in-process + (cache hit, nothing is reloaded) or, in the future, as an independently + enqueued job (refs travel, the heavy objects are derived again). + + Parameters + ---------- + refs : Dict[str, Any], optional + Initial JSON serializable references. + """ + + def __init__(self, refs: Optional[Dict[str, Any]] = None) -> None: + self._refs: Dict[str, Any] = {} + self._cache: Dict[str, Any] = {} + + for key, value in (refs or {}).items(): + self.put_ref(key, value) + + @property + def refs(self) -> Dict[str, Any]: + """A deep copy of the serializable references held by the context.""" + return copy.deepcopy(self._refs) + + def put_ref(self, key: str, value: Any) -> None: + """Store a durable, JSON serializable reference. + + The value is deep-copied on the way in, so a mutable dict handed in + by the caller (e.g. an ORM-attached ``dict`` column) is never aliased. + Without this, a later in-place edit made through the context — such as + ``ModelFactory.update_parameters`` rewriting a nested ``fixed_value`` + — would silently write through to the caller's object. + + Parameters + ---------- + key : str + Name of the reference. + value : Any + Value to store. Must be JSON serializable. + + Raises + ------ + UnitContractError + If the value cannot be serialized to JSON. + """ + try: + json.dumps(value) + except (TypeError, ValueError) as e: + raise UnitContractError( + f"Context reference '{key}' is not JSON serializable " + f"({type(value).__name__}). Only ids, paths and plain data can " + "cross a unit boundary; store live objects with put() instead." + ) from e + + self._refs[key] = copy.deepcopy(value) + + def put(self, key: str, value: Any) -> None: + """Store a live object in the in-process cache. + + Parameters + ---------- + key : str + Name of the value. + value : Any + Any object. It is never serialized. + """ + self._cache[key] = value + + def get(self, key: str, default: Any = None) -> Any: + """Retrieve a value, looking in the cache before the references. + + A live object from the cache is returned by reference — that is the + point of the cache half. A reference is returned as a deep copy, so + mutating what comes back never reaches into the context's own state + or, transitively, into whatever object a ``put_ref`` call was given. + + Parameters + ---------- + key : str + Name of the value. + default : Any, optional + Returned when the key is absent from both halves. + + Returns + ------- + Any + The cached object, a copy of the stored reference, or ``default``. + """ + if key in self._cache: + return self._cache[key] + if key in self._refs: + return copy.deepcopy(self._refs[key]) + return default + + def require(self, key: str) -> Any: + """Retrieve a value, failing when it is absent. + + See :meth:`get` for the copy-on-read guarantee for references. + + Parameters + ---------- + key : str + Name of the value. + + Returns + ------- + Any + The cached object or a copy of the stored reference. + + Raises + ------ + UnitContractError + If the key is present in neither half of the context. + """ + if key in self._cache: + return self._cache[key] + if key in self._refs: + return copy.deepcopy(self._refs[key]) + + raise UnitContractError( + f"Context key '{key}' is not available. " + f"Present keys: {sorted(set(self._cache) | set(self._refs))}." + ) + + def has(self, key: str) -> bool: + """Return whether a key is present in either half of the context.""" + return key in self._cache or key in self._refs + + def clear_cache(self) -> None: + """Drop every live object, keeping the references. + + Called by the orchestrator before ``gc.collect()`` so datasets, splits + and models become collectable as soon as the job is done. + """ + self._cache.clear() + + def to_dict(self) -> Dict[str, Any]: + """Serialize the context to the references that can cross a process. + + Returns + ------- + Dict[str, Any] + A deep copy of the JSON serializable half of the context. + """ + return self.refs + + @classmethod + def from_dict(cls, refs: Dict[str, Any]) -> "ExecutionContext": + """Rebuild a context from previously serialized references. + + Parameters + ---------- + refs : Dict[str, Any] + References as returned by :meth:`to_dict`. + + Returns + ------- + ExecutionContext + A context with an empty cache; live objects are derived on demand. + """ + return cls(refs=refs) diff --git a/DashAI/back/units/evaluate_model_unit.py b/DashAI/back/units/evaluate_model_unit.py new file mode 100644 index 000000000..bdce525f7 --- /dev/null +++ b/DashAI/back/units/evaluate_model_unit.py @@ -0,0 +1,92 @@ +"""Unit that computes and stores the final metrics of a trained model.""" + +import logging + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + list_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Metric +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__) + +DEFAULT_SPLITS = ["TRAIN", "VALIDATION", "TEST"] + + +class EvaluateModelSchema(BaseSchema): + splits: schema_field( + list_field(enum_field(enum=DEFAULT_SPLITS)), + placeholder=DEFAULT_SPLITS, + description=MultilingualString( + en="Data splits the model is evaluated on.", + es="Particiones de datos sobre las que se evalúa el modelo.", + pt="Partições de dados sobre as quais o modelo é avaliado.", + de="Datenteilmengen, auf denen das Modell ausgewertet wird.", + zh="用于评估模型的数据划分。", + ), + alias=MultilingualString( + en="Splits", + es="Particiones", + pt="Partições", + de="Teilmengen", + zh="数据划分", + ), + ) # type: ignore + + +class EvaluateModelUnit(BaseUnit): + """Compute the final metrics of a trained model, once per split. + + The unit is idempotent: a split whose final metrics were already logged + (for instance by a model that evaluates itself while training) is skipped. + + Which metrics are computed is decided by the model, not by this unit: + ``ModelFactory`` attaches the metric classes to the model instance, and + ``BaseModel.calculate_metrics`` is ``final``. That method persists the + rows through a session of its own, so these writes are not part of the + transaction the calling job controls. + """ + + SCHEMA = EvaluateModelSchema + + REQUIRES = ("model", "run_id") + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + model = ctx.require("model") + # ctx.require, not ctx.get: run_id is what the idempotency query below + # filters on. A silently-None run_id would match no existing metric + # row regardless of what was actually logged, and — if the model + # were also somehow detached from its run — calculate_metrics would + # then no-op (base_model.py's ``if not metrics or not self.run_id``), + # so the unit would "succeed" having written nothing. + run_id = ctx.require("run_id") + splits = [SplitEnum[name] for name in self.config.get("splits", DEFAULT_SPLITS)] + + try: + for split in splits: + with session_factory() as db: + already_logged = ( + db.query(Metric) + .filter_by(run_id=run_id, split=split, level=LevelEnum.LAST) + .first() + ) + if already_logged: + continue + + model.calculate_metrics(split=split, level=LevelEnum.LAST) + except Exception as e: + log.exception(e) + raise JobError( + f"Metric calculation failed {e}", + ) from e diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py new file mode 100644 index 000000000..8fc16f958 --- /dev/null +++ b/DashAI/back/units/fit_model_unit.py @@ -0,0 +1,230 @@ +"""Unit that fits a model, optionally searching for its hyperparameters.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_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.optimizers.base_optimizer import BaseOptimizer + +log = logging.getLogger(__name__) + + +class FitModelSchema(BaseSchema): + optimizer: schema_field( + component_field(parent="BaseOptimizer"), + placeholder={"component": "OptunaOptimizer", "params": {}}, + description=MultilingualString( + en="Optimizer used to search for hyperparameters, along with its own " + "configuration. Only used when the model declares optimizable " + "parameters.", + es="Optimizador usado para buscar hiperparámetros, junto con su propia " + "configuración. Solo se usa cuando el modelo declara parámetros " + "optimizables.", + pt="Otimizador usado para procurar hiperparâmetros, junto com a sua " + "própria configuração. Só é usado quando o modelo declara parâmetros " + "otimizáveis.", + de="Optimierer für die Hyperparametersuche samt seiner eigenen " + "Konfiguration. Wird nur verwendet, wenn das Modell optimierbare " + "Parameter deklariert.", + zh="用于搜索超参数的优化器及其自身配置。仅当模型声明了可优化参数时使用。", + ), + alias=MultilingualString( + en="Optimizer", + es="Optimizador", + pt="Otimizador", + de="Optimierer", + zh="优化器", + ), + ) # type: ignore + goal_metric: schema_field( + string_field(), + placeholder="Accuracy", + description=MultilingualString( + en="Metric the hyperparameter search optimizes.", + es="Métrica que optimiza la búsqueda de hiperparámetros.", + pt="Métrica que a procura de hiperparâmetros otimiza.", + de="Metrik, die die Hyperparametersuche optimiert.", + zh="超参数搜索所优化的指标。", + ), + alias=MultilingualString( + en="Goal metric", + es="Métrica objetivo", + pt="Métrica objetivo", + de="Zielmetrik", + zh="目标指标", + ), + ) # type: ignore + + +class FitModelUnit(BaseUnit): + """Train a model, running a hyperparameter search when there is one to run. + + Hyperparameter optimization is a fitting strategy rather than a separate + step: it returns a fitted model, and the trial plots are a by-product only + that branch produces. Both paths therefore live in this unit. + + ``validate`` resolves the optimizer and the goal metric so an impossible + configuration is rejected before the job reports that training started. + + The optimizer is configured as a component field, so its value is + ``{"component": , "params": {...}}`` and the front renders the + chosen optimizer's own form underneath. + """ + + SCHEMA = FitModelSchema + + REQUIRES = ( + "model", + "factory", + "optimizable_parameters", + "model_parameters", + "x", + "y", + "task", + ) + PROVIDES = ("model", "plot_paths") + + def validate(self, ctx: ExecutionContext) -> None: + # ctx.require, not ctx.get: "optimizable_parameters" is one of this + # unit's REQUIRES, so its absence means BuildModelUnit hasn't run yet + # — a call-order mistake, not "there is nothing to optimize". Only an + # empty value (the key present, genuinely no optimizable parameters) + # skips the optimizer/goal-metric checks below, so no registry lookup + # is needed either. + if not ctx.require("optimizable_parameters"): + return + + from kink import di + + component_registry = di["component_registry"] + goal_metric_name: str = self.config["goal_metric"] + optimizer_name: str = self.config["optimizer"]["component"] + + try: + # The whole registry entry, not the class: the optimizer reads + # metadata["maximize"] from it to pick a direction. + goal_metric = component_registry[goal_metric_name] + except Exception as e: + log.exception(e) + raise JobError( + f"Metric is not compatible with the Task. {e}", + ) from e + + try: + optimizer_class = component_registry[optimizer_name]["class"] + optimizer: "BaseOptimizer" = optimizer_class( + **self.config["optimizer"]["params"] + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instantiating optimizer {optimizer_name}, {e}", + ) from e + + ctx.put("goal_metric", goal_metric) + ctx.put("optimizer", optimizer) + + def execute(self, ctx: ExecutionContext) -> None: + import os + import pickle + + from kink import di + + config = di["config"] + + model = ctx.require("model") + x = ctx.require("x") + y = ctx.require("y") + optimizable_parameters = ctx.require("optimizable_parameters") + + plot_paths = [] + try: + if not optimizable_parameters: + model.train(x["train"], y["train"], x["validation"], y["validation"]) + else: + # __call__ always runs validate() immediately before execute(), + # so "optimizer"/"goal_metric" are already in ctx here. + optimizer = ctx.require("optimizer") + goal_metric = ctx.require("goal_metric") + factory = ctx.require("factory") + run_id = ctx.get("run_id") + + optimizer.optimize( + model, + x, + y, + optimizable_parameters, + goal_metric, + ctx.require("task"), + ) + model = optimizer.get_model() + best_params = optimizer.get_best_params() + + self._assert_model_keeps_its_runtime_state(model, run_id) + + # ctx.require already hands back an isolated copy of the + # stored parameter tree, so update_parameters is free to + # mutate it without touching the Run row it came from. + old_parameters = ctx.require("model_parameters") + ctx.put_ref( + "best_parameters", + factory.update_parameters(old_parameters, best_params), + ) + + # Generate hyperparameter plot + from DashAI.back.core.artifacts import normalize_artifacts + + trials = optimizer.get_trials_values() + plot_filenames, plots = optimizer.create_plots( + trials, + run_id, + n_params=len(optimizable_parameters), + goal_metric=goal_metric, + ) + normalized_plots = normalize_artifacts(plots) + for filename, plot in zip( + plot_filenames, normalized_plots, strict=False + ): + plot_path = os.path.join(config["RUNS_PATH"], filename) + with open(plot_path, "wb") as file: + pickle.dump(plot, file) + plot_paths.append(plot_path) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + + ctx.put("model", model) + ctx.put_ref("plot_paths", plot_paths) + + @staticmethod + def _assert_model_keeps_its_runtime_state(model, run_id) -> None: + """Fail loudly if the optimizer returned a model that cannot log metrics. + + ``ModelFactory`` attaches the run id, the data splits and the metric + classes to the model instance, and optimizers are expected to return + that same instance. If one ever returns a fresh object instead, + ``calculate_metrics`` would return early and the run would finish with + no metrics at all instead of failing. + """ + if run_id is None: + return + + if getattr(model, "run_id", None) is None: + raise JobError( + "The optimizer returned a model detached from its run: metrics " + "could not be computed for it. Optimizers must return the same " + "model instance they received." + ) diff --git a/DashAI/back/units/load_dataset_unit.py b/DashAI/back/units/load_dataset_unit.py new file mode 100644 index 000000000..dc3244645 --- /dev/null +++ b/DashAI/back/units/load_dataset_unit.py @@ -0,0 +1,74 @@ +"""Unit that loads a stored dataset into memory.""" + +import logging +from typing import TYPE_CHECKING + +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 Dataset +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 LoadDatasetSchema(BaseSchema): + dataset_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the stored dataset to load.", + es="Identificador del conjunto de datos almacenado a cargar.", + pt="Identificador do conjunto de dados armazenado a carregar.", + de="Kennung des zu ladenden gespeicherten Datensatzes.", + zh="要加载的已存储数据集的标识符。", + ), + alias=MultilingualString( + en="Dataset", + es="Conjunto de datos", + pt="Conjunto de dados", + de="Datensatz", + zh="数据集", + ), + ) # type: ignore + + +class LoadDatasetUnit(BaseUnit): + """Load a dataset from disk into the execution context. + + Resolves the dataset row to find where it is stored and materialises it, so + downstream units receive a dataset instead of an identifier. + """ + + SCHEMA = LoadDatasetSchema + + PROVIDES = ("dataset",) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + session_factory = di["session_factory"] + dataset_id: int = self.config["dataset_id"] + + with session_factory() as db: + dataset: Dataset = db.get(Dataset, dataset_id) + if not dataset: + raise JobError(f"Dataset {dataset_id} does not exist in DB.") + file_path = dataset.file_path + + try: + loaded_dataset: "DashAIDataset" = load_dataset(f"{file_path}/dataset") + except Exception as e: + log.exception(e) + raise JobError( + f"Can not load dataset from path {file_path}", + ) from e + + ctx.put_ref("dataset_id", dataset_id) + ctx.put("dataset", loaded_dataset) diff --git a/DashAI/back/units/prepare_and_split_unit.py b/DashAI/back/units/prepare_and_split_unit.py new file mode 100644 index 000000000..283da4fb1 --- /dev/null +++ b/DashAI/back/units/prepare_and_split_unit.py @@ -0,0 +1,182 @@ +"""Unit that prepares a dataset for a task and splits it into train/val/test.""" + +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 PrepareAndSplitSchema(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: schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + 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="用作模型输入的列名。", + ), + 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 has to predict.", + es="Nombres de las columnas que el modelo debe predecir.", + pt="Nomes das colunas que o modelo deve prever.", + de="Namen der Spalten, die das Modell vorhersagen soll.", + zh="模型需要预测的列名。", + ), + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + ) # type: ignore + splits: schema_field( + dict, + placeholder={ + "splitType": "random", + "train": 0.7, + "test": 0.1, + "validation": 0.2, + }, + description=MultilingualString( + en="Split configuration: a split type plus either train/test/" + "validation index lists or proportions.", + es="Configuración de partición: un tipo de partición y listas de " + "índices o proporciones para entrenamiento/prueba/validación.", + pt="Configuração de divisão: um tipo de divisão e listas de " + "índices ou proporções para treino/teste/validação.", + de="Split-Konfiguration: ein Split-Typ sowie entweder Index-" + "Listen oder Anteile für Training/Test/Validierung.", + zh="划分配置：划分类型，以及训练/测试/验证的索引列表或比例。", + ), + alias=MultilingualString( + en="Splits", + es="Particiones", + pt="Partições", + de="Teilmengen", + zh="数据划分", + ), + ) # type: ignore + + +class PrepareAndSplitUnit(BaseUnit): + """Validate a dataset against a task and split it into train/val/test. + + Runs the task's own validation, counts the labels, applies the requested + split configuration and separates features from targets. + """ + + SCHEMA = PrepareAndSplitSchema + + REQUIRES = ("dataset",) + PROVIDES = ("x", "y", "n_labels", "task", "split_indexes") + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.dataloaders.classes.dashai_dataset import ( + prepare_for_model_session, + select_columns, + split_dataset, + ) + + component_registry = di["component_registry"] + + task_name: str = self.config["task_name"] + input_columns = self.config["input_columns"] + output_columns = self.config["output_columns"] + splits = self.config["splits"] + + loaded_dataset = ctx.require("dataset") + + 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 + + try: + prepared_dataset = task.prepare_for_task( + dataset=loaded_dataset, + input_columns=input_columns, + output_columns=output_columns, + ) + n_labels = task.num_labels(prepared_dataset, output_columns[0]) + + prepared_dataset, splits = prepare_for_model_session( + dataset=prepared_dataset, + splits=splits, + output_columns=output_columns, + ) + + split_indexes = { + "train_indexes": splits["train_indexes"], + "test_indexes": splits["test_indexes"], + "val_indexes": splits["val_indexes"], + } + + x, y = select_columns( + prepared_dataset, + input_columns, + output_columns, + ) + + x = split_dataset(x) + y = split_dataset(y) + + except Exception as e: + log.exception(e) + raise JobError( + f"""Can not prepare Dataset {ctx.get("dataset_id")} + for Task {task_name}""", + ) from e + + ctx.put_ref("task_name", task_name) + ctx.put_ref("split_indexes", split_indexes) + ctx.put("task", task) + ctx.put("n_labels", n_labels) + ctx.put("x", x) + ctx.put("y", y) diff --git a/DashAI/back/units/save_model_unit.py b/DashAI/back/units/save_model_unit.py new file mode 100644 index 000000000..5b2d5e3c0 --- /dev/null +++ b/DashAI/back/units/save_model_unit.py @@ -0,0 +1,41 @@ +"""Unit that persists a trained model to disk.""" + +import logging + +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 SaveModelUnit(BaseUnit): + """Write a trained model under the runs directory, keyed by its run id. + + Takes no configuration: the destination is derived from the run the model + belongs to, so a re-run overwrites its own artifact and never another's. + """ + + REQUIRES = ("model",) + PROVIDES = ("model_path",) + + def execute(self, ctx: ExecutionContext) -> None: + import os + + from kink import di + + config = di["config"] + + model = ctx.require("model") + run_id = ctx.require("run_id") + + try: + model_path = os.path.join(config["RUNS_PATH"], str(run_id)) + model.save(model_path) + except Exception as e: + log.exception(e) + raise JobError( + "Model saving failed", + ) from e + + ctx.put_ref("model_path", model_path) diff --git a/tests/back/api/test_model_job_orchestration.py b/tests/back/api/test_model_job_orchestration.py new file mode 100644 index 000000000..f07429252 --- /dev/null +++ b/tests/back/api/test_model_job_orchestration.py @@ -0,0 +1,372 @@ +"""End to end regression net for the ModelJob orchestration. + +``test_jobs.py`` accepts either ``finished`` or ``error`` as a job outcome, so +it cannot catch a job that silently stops doing part of its work. These tests +pin the observable contract of a successful run: status transitions, the +columns written on ``Run``, the ``Metric`` rows and the saved model artifact. +""" + +import json +import os + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.enums.status import RunStatus +from DashAI.back.core.schema_fields import BaseSchema, int_field, schema_field +from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.job.base_job import JobError +from DashAI.back.job.model_job import ModelJob +from DashAI.back.metrics.base_metric import BaseMetric +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.tasks.base_task import BaseTask + + +class OrchestrationTask(BaseTask): + name: str = "OrchestrationTask" + 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 + + def num_labels(self, dataset, output_column): + return 3 + + +class OrchestrationModel(BaseModel): + """Model that records the data it was trained with.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + + def __init__(self, **kwargs): + self.trained_with = None + + def save(self, filename): + joblib.dump({"trained_with": self.trained_with}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.trained_with = { + "train": x_train.shape[0], + "validation": None if x_validation is None else x_validation.shape[0], + } + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + +class TunableModelSchema(BaseSchema): + n_estimators: schema_field( + int_field(gt=0), + placeholder=2, + description="Number of estimators.", + ) # type: ignore + + +class TunableModel(BaseModel): + """Model with an optimizable parameter, to exercise the search branch.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + SCHEMA = TunableModelSchema + + def __init__(self, n_estimators=2, **kwargs): + self.n_estimators = n_estimators + + def save(self, filename): + joblib.dump({"n_estimators": self.n_estimators}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + +class OrchestrationMetric(BaseMetric): + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.5 + + +class TunableMetric(BaseMetric): + """Metric whose score depends on the hyperparameter being searched.""" + + COMPATIBLE_COMPONENTS = ["OrchestrationTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.25 + + +@pytest.fixture(scope="module", name="orchestration_registry", autouse=True) +def setup_orchestration_registry(client): + container = client.app.container + sentinel = object() + services = container._services + old = services.get("component_registry", sentinel) + + services["component_registry"] = ComponentRegistry( + initial_components=[ + OrchestrationTask, + OrchestrationModel, + TunableModel, + OrchestrationMetric, + TunableMetric, + CSVDataLoader, + ModelJob, + OptunaOptimizer, + ] + ) + yield services["component_registry"] + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session( + client: TestClient, dataset_1: Dataset, orchestration_registry +): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="OrchestrationSession", + task_name="OrchestrationTask", + input_columns=["SepalLengthCm", "SepalWidthCm"], + output_columns=["Species"], + train_metrics=["OrchestrationMetric"], + validation_metrics=["OrchestrationMetric"], + test_metrics=["OrchestrationMetric"], + splits=json.dumps( + { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + "splitType": "random", + } + ), + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + yield model_session.id + + +def _create_run(client: TestClient, model_session_id: int, model_name: str) -> int: + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + model_name=model_name, + parameters={}, + optimizer_name="", + optimizer_parameters={}, + goal_metric="", + name="OrchestrationRun", + ) + db.add(run) + db.commit() + db.refresh(run) + return run.id + + +@pytest.fixture(name="finished_run") +def run_a_successful_job(client: TestClient, model_session_id: int) -> Run: + run_id = _create_run(client, model_session_id, "OrchestrationModel") + ModelJob(run_id=run_id).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + return db.get(Run, run_id) + + +def test_successful_run_reaches_finished(finished_run: Run): + assert finished_run.status == RunStatus.FINISHED + + +def test_successful_run_stamps_its_timestamps(finished_run: Run): + assert finished_run.start_time is not None + assert finished_run.end_time is not None + assert finished_run.end_time >= finished_run.start_time + + +def test_successful_run_persists_the_split_indexes(finished_run: Run): + split_indexes = json.loads(finished_run.split_indexes) + + assert set(split_indexes) == {"train_indexes", "test_indexes", "val_indexes"} + assert len(split_indexes["train_indexes"]) > 0 + assert len(split_indexes["val_indexes"]) > 0 + assert len(split_indexes["test_indexes"]) > 0 + + +def test_successful_run_saves_the_model_artifact(finished_run: Run): + assert finished_run.run_path is not None + assert finished_run.run_path.endswith(str(finished_run.id)) + assert os.path.exists(finished_run.run_path) + + +def test_the_model_was_trained_with_the_train_and_validation_splits( + finished_run: Run, +): + saved = joblib.load(finished_run.run_path) + + assert saved["trained_with"]["train"] > 0 + assert saved["trained_with"]["validation"] > 0 + + +def test_successful_run_writes_a_last_metric_for_every_split( + client: TestClient, finished_run: Run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + metrics = db.query(Metric).filter_by(run_id=finished_run.id).all() + + by_split = {metric.split: metric for metric in metrics} + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST} + for metric in metrics: + assert metric.level == LevelEnum.LAST + assert metric.name == "OrchestrationMetric" + assert metric.value == 0.5 + + +@pytest.fixture(scope="module", name="tuned_run") +def run_a_hyperparameter_search(client: TestClient, model_session_id: int) -> Run: + """Run the hyperparameter search branch end to end. + + This branch had no coverage at all: every other test creates runs with an + empty ``optimizer_name`` and no optimizable parameters, so the whole + optimize/plot path was never executed by the suite. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + model_name="TunableModel", + parameters={ + "n_estimators": { + "optimize": True, + "lower_bound": 1, + "upper_bound": 5, + "fixed_value": 2, + } + }, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 3, + "sampler": "TPESampler", + "pruner": None, + }, + goal_metric="TunableMetric", + name="TunedRun", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + return db.get(Run, run_id) + + +def test_hyperparameter_search_reaches_finished(tuned_run: Run): + assert tuned_run.status == RunStatus.FINISHED + + +def test_hyperparameter_search_persists_the_best_parameters(tuned_run: Run): + best = tuned_run.parameters["n_estimators"]["fixed_value"] + + assert 1 <= best <= 5 + + +def test_hyperparameter_search_saves_two_plots_for_a_single_parameter( + tuned_run: Run, +): + assert tuned_run.plot_history_path is not None + assert tuned_run.plot_slice_path is not None + assert os.path.exists(tuned_run.plot_history_path) + assert os.path.exists(tuned_run.plot_slice_path) + # create_plots only produces the contour and importance plots when more + # than one hyperparameter is being searched. + assert tuned_run.plot_contour_path is None + assert tuned_run.plot_importance_path is None + + +def test_the_tuned_model_still_logs_its_metrics(client: TestClient, tuned_run: Run): + """The optimizer must return the same model instance it received. + + ``ModelFactory`` hangs the run id, the splits and the metric classes off + the model. If an optimizer returned a fresh object instead, + ``calculate_metrics`` would silently return and the run would finish with + no metrics rather than fail. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + metrics = db.query(Metric).filter_by(run_id=tuned_run.id).all() + + last_splits = {m.split for m in metrics if m.level == LevelEnum.LAST} + assert last_splits == {SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST} + + # The optimizer logs a metric per trial while searching. + trial_metrics = [m for m in metrics if m.level == LevelEnum.TRIAL] + assert len(trial_metrics) > 0 + + +def test_a_missing_run_is_reported_instead_of_crashing(client: TestClient): + with pytest.raises(JobError, match="Run 987654 does not exist in DB."): + ModelJob(run_id=987654).run() + + +def test_an_unknown_model_leaves_the_run_in_error( + client: TestClient, model_session_id: int +): + run_id = _create_run(client, model_session_id, "ThereIsNoSuchModel") + + with pytest.raises(JobError, match="Unable to find Model with name"): + ModelJob(run_id=run_id).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.get(Run, run_id).status == RunStatus.ERROR diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py new file mode 100644 index 000000000..87de3d44a --- /dev/null +++ b/tests/back/api/test_units_api.py @@ -0,0 +1,79 @@ +"""The atomic units are exposed as regular registry components.""" + +import pytest +from fastapi.testclient import TestClient + +EXPECTED_UNITS = { + "LoadDatasetUnit", + "PrepareAndSplitUnit", + "BuildModelUnit", + "FitModelUnit", + "EvaluateModelUnit", + "SaveModelUnit", +} + + +@pytest.fixture(name="units", scope="module") +def get_units(client: TestClient): + response = client.get("/api/v1/component/?select_types=Unit") + assert response.status_code == 200, response.text + return {component["name"]: component for component in response.json()} + + +def test_every_unit_is_registered(units): + assert set(units) == EXPECTED_UNITS + + +def test_units_are_registered_under_the_unit_type(units): + for unit in units.values(): + assert unit["type"] == "Unit" + + +def test_units_expose_a_schema_the_front_can_render(units): + for name, unit in units.items(): + assert unit["configurable_object"] is True, name + assert "properties" in unit["schema"], name + + +def test_unit_schemas_describe_their_configuration(units): + assert "dataset_id" in units["LoadDatasetUnit"]["schema"]["properties"] + assert set(units["PrepareAndSplitUnit"]["schema"]["properties"]) == { + "task_name", + "input_columns", + "output_columns", + "splits", + } + assert "model" in units["BuildModelUnit"]["schema"]["properties"] + assert "optimizer" in units["FitModelUnit"]["schema"]["properties"] + + +def test_component_fields_tell_the_front_which_components_to_offer(units): + """The recursive part of the schema system. + + A component field does not inline the chosen component's schema; it + carries a ``parent`` hint so the front can list the candidates and then + fetch that component's own schema to render the nested form. + """ + model = units["BuildModelUnit"]["schema"]["properties"]["model"] + optimizer = units["FitModelUnit"]["schema"]["properties"]["optimizer"] + + assert model["parent"] == "BaseModel" + assert optimizer["parent"] == "BaseOptimizer" + assert set(model["properties"]) == {"component", "params"} + + +def test_a_component_field_parent_resolves_to_real_components(client: TestClient): + response = client.get("/api/v1/component/?component_parent=BaseOptimizer") + assert response.status_code == 200, response.text + + names = {component["name"] for component in response.json()} + assert "OptunaOptimizer" in names + + +def test_units_do_not_leak_into_the_job_listing(client: TestClient): + response = client.get("/api/v1/component/?select_types=Job") + assert response.status_code == 200, response.text + + job_names = {component["name"] for component in response.json()} + assert not (job_names & EXPECTED_UNITS) + assert "ModelJob" in job_names diff --git a/tests/back/conftest.py b/tests/back/conftest.py index e1a96dfb6..a9d525822 100644 --- a/tests/back/conftest.py +++ b/tests/back/conftest.py @@ -7,6 +7,7 @@ import pytest from DashAI.back.dependencies.job_queues.huey_job_queue import HueyJobQueue +from tests.back.scratch import clear_scratch TEST_PATH = pathlib.Path("tmp") TEST_DATASETS_PATH = pathlib.Path("./tests/back/test_datasets") @@ -18,6 +19,14 @@ def test_path(): return TEST_PATH +@pytest.fixture(scope="session", autouse=True) +def _clean_scratch_directories(): + """Drop the dataset caches the suite writes outside the repository.""" + clear_scratch() + yield + clear_scratch() + + @pytest.fixture(scope="session", autouse=True) def random_state(): return RANDOM_STATE diff --git a/tests/back/dataloaders/base_tabular_dataloader_tests.py b/tests/back/dataloaders/base_tabular_dataloader_tests.py index 7a6a43c23..4736b66bd 100644 --- a/tests/back/dataloaders/base_tabular_dataloader_tests.py +++ b/tests/back/dataloaders/base_tabular_dataloader_tests.py @@ -10,6 +10,7 @@ from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset, split_dataset from DashAI.back.dataloaders.classes.dataloader import BaseDataLoader +from tests.back.scratch import scratch_dir # TODO: Test no header, empty file, bad split folder structure. @@ -70,7 +71,7 @@ def _test_load_data_from_file( # load data dataset = dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders", + temp_path=scratch_dir("dataloaders"), params=params, ) @@ -116,7 +117,7 @@ def _test_load_data_from_zip( dataset = dataloder_instance.load_data( filepath_or_buffer=str(dataset_path), - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) dataset = split_dataset(dataset) @@ -171,7 +172,7 @@ def _test_dataloader_with_missing_required_params( ): dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) @@ -205,6 +206,6 @@ def _test_dataloader_try_to_load_a_invalid_datasets( ): dataloder_instance.load_data( filepath_or_buffer=file, - temp_path="tests/back/dataloaders/iris", + temp_path=scratch_dir("dataloaders", "iris"), params=params, ) diff --git a/tests/back/dataloaders/test_dashai_dataset.py b/tests/back/dataloaders/test_dashai_dataset.py index 0aec2c50c..5dbd572b9 100644 --- a/tests/back/dataloaders/test_dashai_dataset.py +++ b/tests/back/dataloaders/test_dashai_dataset.py @@ -23,6 +23,7 @@ update_dataset_splits, validate_inputs_outputs, ) +from tests.back.scratch import scratch_dir from tests.back.test_datasets_generator import CSVTestDatasetGenerator @@ -61,7 +62,7 @@ def load_test_datasetdict(test_datasets_path: pathlib.Path) -> DatasetDict: test_datasetdict = CSVDataLoader().load_data( filepath_or_buffer=file, - temp_path=str(test_datasets_path), + temp_path=scratch_dir("dashai_dataset"), params={"separator": ";"}, ) diff --git a/tests/back/explainers/test_explainers.py b/tests/back/explainers/test_explainers.py index 2c854683a..14a3b8648 100644 --- a/tests/back/explainers/test_explainers.py +++ b/tests/back/explainers/test_explainers.py @@ -24,6 +24,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -46,7 +47,7 @@ def tabular_model_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_lib_explainers.py b/tests/back/explainers/test_lib_explainers.py index acd8cf9de..46b8917cb 100644 --- a/tests/back/explainers/test_lib_explainers.py +++ b/tests/back/explainers/test_lib_explainers.py @@ -22,6 +22,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -44,7 +45,7 @@ def tabular_dataset_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_new_explainers.py b/tests/back/explainers/test_new_explainers.py index 775a59366..c53cf0e49 100644 --- a/tests/back/explainers/test_new_explainers.py +++ b/tests/back/explainers/test_new_explainers.py @@ -21,6 +21,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir INPUT_COLUMNS = [ "SepalLengthCm", @@ -43,7 +44,7 @@ def tabular_model_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/explainers/test_task_explainers.py b/tests/back/explainers/test_task_explainers.py index c656efc59..aba128914 100644 --- a/tests/back/explainers/test_task_explainers.py +++ b/tests/back/explainers/test_task_explainers.py @@ -26,6 +26,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir REGRESSION_INPUT_COLUMNS = [ "SepalLengthCm", @@ -42,7 +43,7 @@ def regression_dataset_fixture(): datasetdict = dataloader.load_data( filepath_or_buffer=dataset_path, - temp_path="tests/back/explainers", + temp_path=scratch_dir("explainers"), params={ "separator": ",", "schema": { diff --git a/tests/back/models/test_bow_text_class_model.py b/tests/back/models/test_bow_text_class_model.py index 1e3dff685..0828c5985 100644 --- a/tests/back/models/test_bow_text_class_model.py +++ b/tests/back/models/test_bow_text_class_model.py @@ -26,6 +26,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(autouse=True, name="test_registry") @@ -58,7 +59,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_deberta_v3_transformer.py b/tests/back/models/test_deberta_v3_transformer.py index a92f4fd5e..fb35b1a21 100644 --- a/tests/back/models/test_deberta_v3_transformer.py +++ b/tests/back/models/test_deberta_v3_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_distilbert_transformer.py b/tests/back/models/test_distilbert_transformer.py index ddf233ef3..60cbf0aa7 100644 --- a/tests/back/models/test_distilbert_transformer.py +++ b/tests/back/models/test_distilbert_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_modernbert_transformer.py b/tests/back/models/test_modernbert_transformer.py index 0801d1ce4..1080a0a50 100644 --- a/tests/back/models/test_modernbert_transformer.py +++ b/tests/back/models/test_modernbert_transformer.py @@ -13,6 +13,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Text +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="splited_dataset") @@ -22,7 +23,7 @@ def splited_dataset_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "data_key": "data", "schema": { diff --git a/tests/back/models/test_tabular_class_models.py b/tests/back/models/test_tabular_class_models.py index 3d99798d0..c5bdf537e 100644 --- a/tests/back/models/test_tabular_class_models.py +++ b/tests/back/models/test_tabular_class_models.py @@ -34,6 +34,7 @@ from DashAI.back.types.categorical import Categorical from DashAI.back.types.utils import save_types_in_arrow_metadata from DashAI.back.types.value_types import Float +from tests.back.scratch import scratch_dir @pytest.fixture(scope="module", name="divided_dataset") @@ -43,7 +44,7 @@ def tabular_model_fixture(): datasetdict = dataloader_test.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/models", + temp_path=scratch_dir("models"), params={ "separator": ",", "schema": { diff --git a/tests/back/scratch.py b/tests/back/scratch.py new file mode 100644 index 000000000..747564b43 --- /dev/null +++ b/tests/back/scratch.py @@ -0,0 +1,50 @@ +"""Scratch directories for tests that load datasets. + +Dataloaders take a ``temp_path`` and hand it straight to HuggingFace as +``cache_dir`` (see ``CSVDataLoader.load_data``), so every test that loads a +dataset leaves an arrow cache and a ``dataset_info.json`` behind. Pointing that +path at the system temp directory keeps those artifacts out of the repository +tree instead of scattering them across ``tests/back/``. +""" + +import os +import pathlib +import shutil +import tempfile + +# Scoped by pid so two pytest processes running at once (two terminals, a CI +# matrix on one machine) never share a root: without this, one session's +# start-of-run ``clear_scratch()`` would rmtree the arrow cache another, +# still-running session has memory mapped. +SCRATCH_ROOT = ( + pathlib.Path(tempfile.gettempdir()) / f"dashai-test-scratch-{os.getpid()}" +) + + +def scratch_dir(*parts: str) -> str: + """Return a scratch directory, creating it if it does not exist. + + Parameters + ---------- + *parts : str + Path segments below the scratch root. Use them to keep unrelated test + modules from sharing a cache, the way the old in-repo paths did. + + Returns + ------- + str + Absolute path to the directory. + """ + path = SCRATCH_ROOT.joinpath(*parts) + path.mkdir(parents=True, exist_ok=True) + return str(path) + + +def clear_scratch() -> None: + """Remove the scratch root, ignoring files still held open. + + Arrow caches stay memory mapped on Windows, so a best effort delete is the + most that can be promised here; the point is that whatever survives lives + outside the repository. + """ + shutil.rmtree(SCRATCH_ROOT, ignore_errors=True) diff --git a/tests/back/tasks/test_tasks.py b/tests/back/tasks/test_tasks.py index 7bce6b40b..618547ec1 100644 --- a/tests/back/tasks/test_tasks.py +++ b/tests/back/tasks/test_tasks.py @@ -20,6 +20,7 @@ from DashAI.back.tasks.text_to_image_generation_task import TextToImageGenerationTask from DashAI.back.tasks.text_to_text_generation_task import TextToTextGenerationTask from DashAI.back.tasks.translation_task import TranslationTask +from tests.back.scratch import scratch_dir def load_csv_into_datasetdict_iris(file_name): @@ -28,7 +29,7 @@ def load_csv_into_datasetdict_iris(file_name): datasetdict = csv_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"separator": ","}, ) schema = { @@ -48,7 +49,7 @@ def load_csv_into_datasetdict_iris_extra(file_name): datasetdict = csv_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"separator": ","}, ) schema = { @@ -148,7 +149,7 @@ def text_classification_dataset_fixture(): dataset = json_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={ "data_key": "data", }, @@ -209,7 +210,7 @@ def translation_dataset_fixture(): dataset = json_dataloader.load_data( filepath_or_buffer=test_dataset_path, - temp_path="tests/back/tasks", + temp_path=scratch_dir("tasks"), params={"data_key": "data"}, ) @@ -271,8 +272,7 @@ def sample_image_fixture(): @pytest.fixture(scope="module", name="temp_path") def temp_path_fixture(): - temp_path = pathlib.Path("tests") / "back" / "tasks" / "temp" - os.makedirs(temp_path, exist_ok=True) + temp_path = pathlib.Path(scratch_dir("tasks", "temp")) yield temp_path # Cleanup after all tests in the module using this fixture have finished if temp_path.exists() and temp_path.is_dir(): diff --git a/tests/back/units/__init__.py b/tests/back/units/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/units/test_base_unit.py b/tests/back/units/test_base_unit.py new file mode 100644 index 000000000..bba78d84d --- /dev/null +++ b/tests/back/units/test_base_unit.py @@ -0,0 +1,132 @@ +"""Tests for the atomic unit base class and its registration.""" + +import pytest + +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +class DummyUnit(BaseUnit): + """Unit that copies a required key into a provided one.""" + + REQUIRES = ("dataset",) + PROVIDES = ("x",) + + def execute(self, ctx: ExecutionContext) -> None: + ctx.put("x", f"prepared {ctx.require('dataset')}") + + +class LyingUnit(BaseUnit): + """Unit that declares an output it never writes.""" + + PROVIDES = ("x",) + + def execute(self, ctx: ExecutionContext) -> None: + pass + + +class ValidatingUnit(BaseUnit): + """Unit that rejects its configuration before running.""" + + def __init__(self, **config) -> None: + super().__init__(**config) + self.validate_calls = 0 + + def validate(self, ctx: ExecutionContext) -> None: + self.validate_calls += 1 + if not self.config.get("allowed", True): + raise ValueError("not allowed") + + def execute(self, ctx: ExecutionContext) -> None: + ctx.put("ran", True) + + +def test_unit_stores_its_configuration(): + unit = DummyUnit(dataset_id=3, columns=["a"]) + + assert unit.config == {"dataset_id": 3, "columns": ["a"]} + + +def test_unit_runs_and_writes_its_output(): + ctx = ExecutionContext() + ctx.put("dataset", "a dataset") + + DummyUnit()(ctx) + + assert ctx.require("x") == "prepared a dataset" + + +def test_unit_refuses_to_run_without_its_required_keys(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + DummyUnit()(ctx) + + +def test_unit_fails_when_it_does_not_deliver_what_it_promised(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'x' is not available"): + LyingUnit()(ctx) + + +def test_validate_is_a_noop_by_default(): + DummyUnit().validate(ExecutionContext()) + + +def test_validate_can_reject_a_configuration_without_executing(): + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=False) + + with pytest.raises(ValueError, match="not allowed"): + unit.validate(ctx) + + assert not ctx.has("ran") + + +def test_call_invokes_validate_automatically_before_execute(): + """Regression: ``unit(ctx)`` alone must run validate(), not just execute(). + + Before this fix a unit's ``validate`` (e.g. ``BuildModelUnit``'s download + gate) only ran if the caller remembered to invoke it separately — an easy + step to forget for any future caller that just does ``unit(ctx)``, the one + call the base class marks as the sanctioned entry point. + """ + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=True) + + unit(ctx) + + assert unit.validate_calls == 1 + assert ctx.require("ran") is True + + +def test_call_stops_at_validate_and_never_reaches_execute(): + ctx = ExecutionContext() + unit = ValidatingUnit(allowed=False) + + with pytest.raises(ValueError, match="not allowed"): + unit(ctx) + + assert not ctx.has("ran") + + +def test_units_register_under_their_own_registry_type(): + """The registry derives TYPE by walking the MRO for a single "Base" + ancestor declaring it. BaseUnit must be that single candidate.""" + registry = ComponentRegistry(initial_components=[DummyUnit]) + + assert registry["DummyUnit"]["type"] == "Unit" + assert [c["name"] for c in registry.get_components_by_types(select="Unit")] == [ + "DummyUnit" + ] + + +def test_registered_units_are_configurable_objects_with_a_schema(): + registry = ComponentRegistry(initial_components=[DummyUnit]) + + component = registry["DummyUnit"] + + assert component["configurable_object"] is True + assert component["schema"] is not None diff --git a/tests/back/units/test_build_model_unit.py b/tests/back/units/test_build_model_unit.py new file mode 100644 index 000000000..8a4c3f624 --- /dev/null +++ b/tests/back/units/test_build_model_unit.py @@ -0,0 +1,74 @@ +"""Tests for BuildModelUnit's model-class resolution and its caching.""" + +import pytest +from kink import di + +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.context import ExecutionContext + + +class _ModelA: + def __init__(self, **kwargs): + pass + + +class _ModelB: + def __init__(self, **kwargs): + pass + + +@pytest.fixture +def fake_registry(): + """A minimal dict-shaped registry: enough for name -> class lookups.""" + registry = { + "ModelA": {"class": _ModelA}, + "ModelB": {"class": _ModelB}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _build_unit(model_name): + return BuildModelUnit( + model={"component": model_name, "params": {}}, + train_metrics=[], + validation_metrics=[], + test_metrics=[], + ) + + +def test_two_build_model_units_in_one_context_resolve_independently(fake_registry): + """Regression: the model-class cache must not be keyed by the context. + + Two ``BuildModelUnit`` instances configured for different models, run + against the same context, must each build their own model. Before this + fix the class was memoized under a context-global ``"model_class"`` key, + so the second unit would silently reuse the first one's resolved class — + and its download-gate check would validate the wrong model too. + """ + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + + a = _build_unit("ModelA") + a(ctx) + model_a = ctx.get("model") + + b = _build_unit("ModelB") + b(ctx) + model_b = ctx.get("model") + + assert isinstance(model_a, _ModelA) + assert isinstance(model_b, _ModelB) + + +def test_resolve_model_class_is_memoized_per_instance_not_shared(fake_registry): + a = _build_unit("ModelA") + b = _build_unit("ModelB") + + assert a._resolve_model_class() is _ModelA + assert b._resolve_model_class() is _ModelB + # Calling again must return the same, still-correct class from the cache. + assert a._resolve_model_class() is _ModelA diff --git a/tests/back/units/test_context.py b/tests/back/units/test_context.py new file mode 100644 index 000000000..4aa201240 --- /dev/null +++ b/tests/back/units/test_context.py @@ -0,0 +1,180 @@ +"""Tests for the execution context shared between atomic units.""" + +import pytest + +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +def test_put_ref_stores_a_serializable_value(): + ctx = ExecutionContext() + ctx.put_ref("run_id", 5) + + assert ctx.get("run_id") == 5 + assert ctx.refs == {"run_id": 5} + + +def test_put_ref_rejects_a_non_serializable_value(): + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="not JSON serializable"): + ctx.put_ref("model", object()) + + +def test_constructor_validates_the_initial_refs(): + with pytest.raises(UnitContractError, match="not JSON serializable"): + ExecutionContext(refs={"model": object()}) + + +def test_put_accepts_a_live_object(): + ctx = ExecutionContext() + sentinel = object() + ctx.put("model", sentinel) + + assert ctx.get("model") is sentinel + assert ctx.refs == {} + + +def test_get_prefers_the_cache_over_the_refs(): + ctx = ExecutionContext(refs={"model_path": "/on/disk"}) + ctx.put("model_path", "/in/memory") + + assert ctx.get("model_path") == "/in/memory" + + +def test_get_falls_back_to_the_refs_and_then_the_default(): + ctx = ExecutionContext(refs={"dataset_id": 3}) + + assert ctx.get("dataset_id") == 3 + assert ctx.get("missing") is None + assert ctx.get("missing", "fallback") == "fallback" + + +def test_require_returns_the_value_when_present(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + assert ctx.require("run_id") == 7 + assert ctx.require("model") == "a model" + + +def test_require_raises_and_lists_the_present_keys(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + with pytest.raises(UnitContractError) as exc_info: + ctx.require("x") + + message = str(exc_info.value) + assert "'x' is not available" in message + assert "'model'" in message + assert "'run_id'" in message + + +def test_has_checks_both_halves(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + assert ctx.has("run_id") + assert ctx.has("model") + assert not ctx.has("x") + + +def test_clear_cache_drops_live_objects_and_keeps_refs(): + ctx = ExecutionContext(refs={"run_id": 7}) + ctx.put("model", "a model") + + ctx.clear_cache() + + assert not ctx.has("model") + assert ctx.require("run_id") == 7 + + +def test_refs_property_returns_a_copy(): + ctx = ExecutionContext(refs={"run_id": 7}) + + ctx.refs["run_id"] = 99 + + assert ctx.require("run_id") == 7 + + +def test_round_trip_through_to_dict_keeps_refs_and_drops_the_cache(): + ctx = ExecutionContext( + refs={"run_id": 7, "split_indexes": {"train_indexes": [0, 1]}} + ) + ctx.put("model", object()) + + restored = ExecutionContext.from_dict(ctx.to_dict()) + + assert restored.require("run_id") == 7 + assert restored.require("split_indexes") == {"train_indexes": [0, 1]} + assert not restored.has("model") + + +def test_to_dict_is_detached_from_the_context(): + ctx = ExecutionContext(refs={"run_id": 7}) + + serialized = ctx.to_dict() + serialized["run_id"] = 99 + + assert ctx.require("run_id") == 7 + + +def test_put_ref_isolates_from_later_mutation_of_the_original_object(): + """Regression: ``put_ref`` must not alias a caller's mutable dict. + + ``BuildModelUnit`` passes ``run.parameters`` — the dict SQLAlchemy attaches + to the ``Run`` row — straight into ``put_ref``. Without a defensive copy, + a later in-place edit reachable through the context (e.g. + ``ModelFactory.update_parameters`` rewriting a nested ``fixed_value`` + during hyperparameter search) would silently write through into the ORM + object, corrupting ``run.parameters`` before it was ever meant to change. + """ + original = {"n_estimators": {"fixed_value": 2}} + ctx = ExecutionContext() + ctx.put_ref("model_parameters", original) + + original["n_estimators"]["fixed_value"] = 999 + + assert ctx.require("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_require_returns_an_isolated_copy_of_a_reference(): + ctx = ExecutionContext( + refs={"model_parameters": {"n_estimators": {"fixed_value": 2}}} + ) + + fetched = ctx.require("model_parameters") + fetched["n_estimators"]["fixed_value"] = 999 + + assert ctx.require("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_get_returns_an_isolated_copy_of_a_reference(): + ctx = ExecutionContext( + refs={"model_parameters": {"n_estimators": {"fixed_value": 2}}} + ) + + fetched = ctx.get("model_parameters") + fetched["n_estimators"]["fixed_value"] = 999 + + assert ctx.get("model_parameters")["n_estimators"]["fixed_value"] == 2 + + +def test_cached_live_objects_are_still_returned_by_reference(): + """The copy-on-read fix must only apply to refs, never to the cache. + + A model, an in-memory dataset, an open session — these have to be the + same live object on every ``get``/``require``. Deep-copying them would + silently break training (the model you fit is not the model you saved) + while looking, from the outside, like nothing went wrong. + """ + + class LiveThing: + pass + + live_object = LiveThing() + ctx = ExecutionContext() + ctx.put("model", live_object) + + assert ctx.get("model") is live_object + assert ctx.require("model") is live_object diff --git a/tests/back/units/test_evaluate_model_unit.py b/tests/back/units/test_evaluate_model_unit.py new file mode 100644 index 000000000..2f6f8c187 --- /dev/null +++ b/tests/back/units/test_evaluate_model_unit.py @@ -0,0 +1,24 @@ +"""Tests for EvaluateModelUnit's contract, independent of a real database.""" + +import pytest + +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit + + +def test_call_refuses_to_run_without_a_run_id(): + """Regression: a missing run_id must fail loudly, not silently no-op. + + The idempotency check filters an existing-metric query by ``run_id``; a + silently-``None`` value would match no row regardless of what was already + logged, and — if the model were also somehow detached from its run — + ``BaseModel.calculate_metrics`` no-ops on a falsy ``run_id`` too. Before + this fix ``run_id`` was read with ``ctx.get`` and absent from + ``REQUIRES``, so the unit could "succeed" having written zero metrics + instead of surfacing the missing wiring. + """ + ctx = ExecutionContext() + ctx.put("model", object()) + + with pytest.raises(UnitContractError, match="'run_id'"): + EvaluateModelUnit()(ctx) diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py new file mode 100644 index 000000000..b4a0bab33 --- /dev/null +++ b/tests/back/units/test_fit_model_unit.py @@ -0,0 +1,40 @@ +"""Tests for FitModelUnit's validation, independent of an actual training run.""" + +import pytest + +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_model_unit import FitModelUnit + + +def _unit(optimizer_name="OptunaOptimizer", goal_metric="Accuracy"): + return FitModelUnit( + optimizer={"component": optimizer_name, "params": {}}, + goal_metric=goal_metric, + ) + + +def test_validate_raises_when_called_before_build_model_unit_has_run(): + """Regression: a missing key must not read as "nothing to optimize". + + ``optimizable_parameters`` is only absent from the context when + ``BuildModelUnit`` hasn't run yet — a call-order mistake, not a model with + no optimizable parameters (that case has the key present but empty). + Before this fix, ``validate`` used ``ctx.get`` and treated both the same, + silently skipping the optimizer/goal-metric checks it exists to run. + """ + ctx = ExecutionContext() + + with pytest.raises(UnitContractError, match="'optimizable_parameters'"): + _unit().validate(ctx) + + +def test_validate_is_a_noop_when_there_are_genuinely_no_optimizable_parameters(): + ctx = ExecutionContext() + ctx.put("optimizable_parameters", []) + + # Should not raise, and should not need the optimizer/goal_metric to + # resolve in the registry. + _unit(optimizer_name="DoesNotExist", goal_metric="DoesNotExist").validate(ctx) + + assert not ctx.has("optimizer") + assert not ctx.has("goal_metric") From fdfee4c2a7ec2b1db6ec258e712a65cff0297313 Mon Sep 17 00:00:00 2001 From: Felipe Date: Sat, 1 Aug 2026 19:05:46 -0400 Subject: [PATCH 02/28] Update error message formatting in PrepareAndSplitUnit and add run_id requirement in SaveModelUnit --- DashAI/back/units/prepare_and_split_unit.py | 3 +-- DashAI/back/units/save_model_unit.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/DashAI/back/units/prepare_and_split_unit.py b/DashAI/back/units/prepare_and_split_unit.py index 283da4fb1..adfaf5928 100644 --- a/DashAI/back/units/prepare_and_split_unit.py +++ b/DashAI/back/units/prepare_and_split_unit.py @@ -170,8 +170,7 @@ def execute(self, ctx: ExecutionContext) -> None: except Exception as e: log.exception(e) raise JobError( - f"""Can not prepare Dataset {ctx.get("dataset_id")} - for Task {task_name}""", + f"Can not prepare Dataset {ctx.get('dataset_id')} for Task {task_name}", ) from e ctx.put_ref("task_name", task_name) diff --git a/DashAI/back/units/save_model_unit.py b/DashAI/back/units/save_model_unit.py index 5b2d5e3c0..22d629b12 100644 --- a/DashAI/back/units/save_model_unit.py +++ b/DashAI/back/units/save_model_unit.py @@ -16,7 +16,7 @@ class SaveModelUnit(BaseUnit): belongs to, so a re-run overwrites its own artifact and never another's. """ - REQUIRES = ("model",) + REQUIRES = ("model", "run_id") PROVIDES = ("model_path",) def execute(self, ctx: ExecutionContext) -> None: From 812dfe3e4dd29c3dc1c6c9ca1a2763fa96e9b53a Mon Sep 17 00:00:00 2001 From: Felipe Date: Sun, 2 Aug 2026 14:30:16 -0400 Subject: [PATCH 03/28] feat: Enhance dataset handling and add SaveDatasetUnit - Updated PrepareAndSplitUnit to require dataset_id and provide task_name. - Introduced SaveDatasetUnit for persisting datasets to disk. - Added comprehensive tests for SaveDatasetUnit and LoadDatasetUnit to ensure correct functionality. - Implemented contract tests for ApplyConverterUnit to validate context handling and converter behavior. - Enhanced unit contract tests to ensure all context keys are declared and properly managed. --- DashAI/back/initial_components.py | 4 + DashAI/back/job/converter_job.py | 312 ++------------ DashAI/back/units/apply_converter_unit.py | 374 +++++++++++++++++ DashAI/back/units/build_model_unit.py | 15 +- DashAI/back/units/fit_model_unit.py | 52 ++- DashAI/back/units/load_dataset_unit.py | 102 ++++- DashAI/back/units/prepare_and_split_unit.py | 10 +- DashAI/back/units/save_dataset_unit.py | 34 ++ tests/back/api/test_converter_job.py | 314 ++++++++++++++ tests/back/api/test_units_api.py | 17 +- tests/back/units/test_apply_converter_unit.py | 388 ++++++++++++++++++ tests/back/units/test_build_model_unit.py | 2 + tests/back/units/test_fit_model_unit.py | 40 +- tests/back/units/test_load_dataset_unit.py | 151 +++++++ tests/back/units/test_save_dataset_unit.py | 64 +++ tests/back/units/test_unit_contracts.py | 135 ++++++ 16 files changed, 1686 insertions(+), 328 deletions(-) create mode 100644 DashAI/back/units/apply_converter_unit.py create mode 100644 DashAI/back/units/save_dataset_unit.py create mode 100644 tests/back/api/test_converter_job.py create mode 100644 tests/back/units/test_apply_converter_unit.py create mode 100644 tests/back/units/test_load_dataset_unit.py create mode 100644 tests/back/units/test_save_dataset_unit.py create mode 100644 tests/back/units/test_unit_contracts.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 1882857ff..5cbcb849a 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -347,11 +347,13 @@ from DashAI.back.tasks.translation_task import TranslationTask # Units +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit 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.load_dataset_unit import LoadDatasetUnit from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit from DashAI.back.units.save_model_unit import SaveModelUnit logging.basicConfig(level=logging.DEBUG) @@ -523,6 +525,8 @@ def get_initial_components(): FitModelUnit, EvaluateModelUnit, SaveModelUnit, + ApplyConverterUnit, + SaveDatasetUnit, # Explainers ContrastiveShap, DiceCounterfactual, diff --git a/DashAI/back/job/converter_job.py b/DashAI/back/job/converter_job.py index fcbfc4d68..f4b8e187b 100644 --- a/DashAI/back/job/converter_job.py +++ b/DashAI/back/job/converter_job.py @@ -1,123 +1,24 @@ import logging -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc -from DashAI.back.api.api_v1.schemas.converter_params import ConverterParams from DashAI.back.dependencies.database.models import Converter from DashAI.back.dependencies.database.models import Dataset as DatasetModel from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker - from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset - logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) -def _rebuild_dataset_with_transformed_columns( - base: "DashAIDataset", - transformed: "DashAIDataset", - scope_column_names: List[str], - scope_column_indexes: List[int], -) -> "DashAIDataset": - """ - Replaces specific columns in the base dataset with columns from the transformed - dataset, preserving their original positions. Also appends any additional columns - that were generated by the transformer at the end. Keeps the features and metadata - consistent. - - Parameters - ---------- - base : DashAIDataset - The original dataset before transformation. - - transformed : DashAIDataset - The dataset resulting from applying a transformer, containing updated and/or - new columns. - - scope_column_names : List[str] - Names of the columns that were originally selected for transformation. - - scope_column_indexes : List[int] - The indices of the columns in the base dataset that were replaced. - Must match the order of scope_column_names. - - Returns - ------- - DashAIDataset - A new dataset with the specified columns replaced in place, new columns - appended, and original metadata and split information preserved. - """ - from DashAI.back.dataloaders.classes.dashai_dataset import modify_table - - original_columns = base.column_names - transformed_cols = transformed.column_names - - transformed_cols_set = set(transformed_cols) - scope_column_names_set = set(scope_column_names) - - removed_cols = [ - col for col in scope_column_names if col not in transformed_cols_set - ] - replacement_cols = [ - col for col in scope_column_names if col in transformed_cols_set - ] - new_cols = [col for col in transformed_cols if col not in scope_column_names_set] - - removed_cols_set = set(removed_cols) - - new_columns_order = [] - seen_cols = set() - for col in original_columns: - if col in removed_cols_set: - continue - if col not in seen_cols: - new_columns_order.append(col) - seen_cols.add(col) - - col_name_mapping = {} - for col in new_cols: - unique_col = col - counter = 1 - while unique_col in seen_cols: - unique_col = f"{col}_{counter}" - counter += 1 - new_columns_order.append(unique_col) - seen_cols.add(unique_col) - col_name_mapping[col] = unique_col - - updated_arrays = {} - for col in replacement_cols: - if col in transformed_cols_set: - updated_arrays[col] = transformed.arrow_table[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed_cols_set: - updated_arrays[unique_col] = transformed.arrow_table[col] - - updated_types = base.types.copy() - - for col in removed_cols: - if col in updated_types: - del updated_types[col] - - for col in replacement_cols: - if col in transformed.types: - updated_types[col] = transformed.types[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed.types: - updated_types[unique_col] = transformed.types[col] - - # Use existing modify_table (imported at module level) - modified_dataset = modify_table(base, updated_arrays, types=updated_types) - modified_dataset = modified_dataset.select_columns(new_columns_order) - - return modified_dataset - - class ConverterJob(BaseJob): """ConverterJob class to modify a dataset by applying a sequence of converters.""" @@ -198,31 +99,9 @@ def run( ) -> None: from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - ) - session_factory = di["session_factory"] - component_registry = di["component_registry"] - - def instantiate_converters( - converter_name: str, - converter_params: ConverterParams, - ) -> object: - # Import the converter - try: - converter_constructor = component_registry[converter_name]["class"] - except KeyError as e: - log.exception(e) - raise JobError( - f"Error importing converter {converter_name}: {e}" - ) from e - - # Get parameters or empty dict if none - converter_parameters = converter_params.get("params", {}) - return converter_constructor(**converter_parameters) + ctx = ExecutionContext() # Extract job parameters converter_id = self.kwargs["converter_id"] @@ -246,21 +125,17 @@ def instantiate_converters( # Get dataset try: + notebook_id = converter.notebook_id dataset_id = converter.notebook.dataset_id dataset = db.get(DatasetModel, dataset_id) - # dataset to edit - dataset_path = f"{converter.notebook.file_path}/dataset" - loaded_dataset = load_dataset(dataset_path) - params = converter.parameters or {} - target_column_index = ( - params["target"].get("idx") - if params.get("target") is not None - else None - ) + # dataset to edit: the notebook's own working copy + LoadDatasetUnit(notebook_id=notebook_id)(ctx) + dataset_path = ctx.require("dataset_path") - if not loaded_dataset: - raise JobError(f"Dataset with path {dataset_path} not found") + # How the converter configuration is stored on the row, not + # part of the transformation itself. + params = converter.parameters or {} except exc.SQLAlchemyError as e: log.exception(e) @@ -268,16 +143,20 @@ def instantiate_converters( db.commit() raise JobError("Error loading dataset info") from e - # Load dataset + apply_converter = ApplyConverterUnit( + converter={ + "component": converter.converter, + "params": params.get("params") or {}, + }, + scope=params.get("scope"), + target=params.get("target"), + ) + + # Validating before the work starts keeps an impossible target + # index reported as a dataset problem, which is where it was + # reported before the job was split into units. try: - # Validate target column index - if target_column_index is not None and ( - int(target_column_index) < 1 - or int(target_column_index) > len(loaded_dataset.features) - ): - raise JobError( - f"Target column index {target_column_index} is out of bounds" - ) + apply_converter.validate(ctx) except Exception as e: log.exception(e) converter.set_status_as_error() @@ -285,142 +164,11 @@ def instantiate_converters( raise JobError(f"Cannot load dataset from {dataset_path}") from e try: - # Get stored converter configurations - converters_stored_info = {converter.converter: converter.parameters} - dataset_original_columns = loaded_dataset.column_names - - # Sort converters by order - converters_sorted_list = sorted( - converters_stored_info.items(), key=lambda x: x[1]["order"] - ) - - i = 0 - converter_instances = [] - - while i < len(converters_sorted_list): - converter_name = converters_sorted_list[i][0] - converter_params = converters_sorted_list[i][1] - # Regular converter - converter_instance = instantiate_converters( - converter_name, - converter_params, - ) - - # Get scope or use default - scope = converter_params.get("scope", {"columns": [], "rows": []}) - - # Add to instances - converter_instances.append( - { - "name": converter_name, - "instance": converter_instance, - "scope": scope, - } - ) - i += 1 - - # Apply each converter in sequence - total_converters = len(converter_instances) - for converter_index, converter_info in enumerate(converter_instances): - converter_instance = converter_info["instance"] - converter_name = converter_info["name"] - converter_scope = converter_info["scope"] - - # Map converter progress onto the 0.2-0.9 band. - self.report_progress( - 0.2 + 0.7 * (converter_index / max(total_converters, 1)), - f"Applying {converter_name}", - ) - log.info(f"Applying converter: {converter_name}") - - columns_scope = [ - column["idx"] - 1 for column in converter_scope["columns"] - ] - scope_column_indexes = sorted(set(columns_scope)) - - if not scope_column_indexes: - scope_column_indexes = list(range(len(loaded_dataset.features))) - - scope_column_names = [ - dataset_original_columns[index] - for index in scope_column_indexes - ] - - rows_scope = [row - 1 for row in converter_scope["rows"]] - scope_rows_indexes = sorted(set(rows_scope)) - - y_dataset_fit = None - target_column_name = None - y_full_transform = None - if target_column_index is not None: - target_column_index_0based = int(target_column_index) - 1 - target_column_name = dataset_original_columns[ - target_column_index_0based - ] - y_dataset_fit = loaded_dataset.select_columns( - [target_column_name] - ) - if scope_rows_indexes: - y_dataset_fit = y_dataset_fit.select(scope_rows_indexes) - y_full_transform = loaded_dataset.select_columns( - [target_column_name] - ) - else: - y_full_transform = y_dataset_fit - - X_dataset_fit = loaded_dataset.select_columns(scope_column_names) - - if scope_rows_indexes: - X_dataset_fit = X_dataset_fit.select(scope_rows_indexes) - - try: - converter_instance = converter_instance.fit( - X_dataset_fit, y_dataset_fit - ) - except ValueError as e: - log.error(f"Validation error in {converter_name}: {e}") - raise JobError( - f"Validation error fitting {converter_name}: {e}" - ) from e - except Exception as e: - log.exception(e) - raise JobError( - f"Error fitting converter {converter_name}: {e}" - ) from e - - if scope_rows_indexes: - X_full_transform = loaded_dataset.select_columns( - scope_column_names - ) - else: - # Same reuse as above: no row-level fit scope means - # X_dataset_fit already covers the full transform scope. - X_full_transform = X_dataset_fit - - try: - transformed_dataset = converter_instance.transform( - X_full_transform, y_full_transform - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error transforming data with {converter_name}: {e}" - ) from e - - if type(converter_instance).CHANGES_ROW_COUNT: - loaded_dataset = transformed_dataset - else: - loaded_dataset = _rebuild_dataset_with_transformed_columns( - loaded_dataset, - transformed_dataset, - scope_column_names, - scope_column_indexes, - ) - - dataset_original_columns = loaded_dataset.column_names + self.report_progress(0.2, f"Applying {converter.converter}") + apply_converter(ctx) self.report_progress(0.95, "Saving dataset") - save_dataset(loaded_dataset, f"{dataset_path}") + SaveDatasetUnit()(ctx) converter.set_status_as_finished() db.commit() db.refresh(dataset) @@ -432,3 +180,5 @@ def instantiate_converters( raise JobError( f"Error applying converters to dataset {dataset_id}: {e}" ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/units/apply_converter_unit.py b/DashAI/back/units/apply_converter_unit.py new file mode 100644 index 000000000..ab4c75998 --- /dev/null +++ b/DashAI/back/units/apply_converter_unit.py @@ -0,0 +1,374 @@ +"""Unit that applies a single converter to the dataset in the context.""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + none_type, + 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 + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + +EMPTY_SCOPE: Dict[str, List] = {"columns": [], "rows": []} + + +def rebuild_dataset_with_transformed_columns( + base: "DashAIDataset", + transformed: "DashAIDataset", + scope_column_names: List[str], +) -> "DashAIDataset": + """ + Replaces specific columns in the base dataset with columns from the transformed + dataset, preserving their original positions. Also appends any additional columns + that were generated by the transformer at the end. Keeps the features and metadata + consistent. + + Parameters + ---------- + base : DashAIDataset + The original dataset before transformation. + + transformed : DashAIDataset + The dataset resulting from applying a transformer, containing updated and/or + new columns. + + scope_column_names : List[str] + Names of the columns that were originally selected for transformation. + + Returns + ------- + DashAIDataset + A new dataset with the specified columns replaced in place, new columns + appended, and original metadata and split information preserved. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import modify_table + + original_columns = base.column_names + transformed_cols = transformed.column_names + + transformed_cols_set = set(transformed_cols) + scope_column_names_set = set(scope_column_names) + + removed_cols = [ + col for col in scope_column_names if col not in transformed_cols_set + ] + replacement_cols = [ + col for col in scope_column_names if col in transformed_cols_set + ] + new_cols = [col for col in transformed_cols if col not in scope_column_names_set] + + removed_cols_set = set(removed_cols) + + new_columns_order = [] + seen_cols = set() + for col in original_columns: + if col in removed_cols_set: + continue + if col not in seen_cols: + new_columns_order.append(col) + seen_cols.add(col) + + col_name_mapping = {} + for col in new_cols: + unique_col = col + counter = 1 + while unique_col in seen_cols: + unique_col = f"{col}_{counter}" + counter += 1 + new_columns_order.append(unique_col) + seen_cols.add(unique_col) + col_name_mapping[col] = unique_col + + updated_arrays = {} + for col in replacement_cols: + if col in transformed_cols_set: + updated_arrays[col] = transformed.arrow_table[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed_cols_set: + updated_arrays[unique_col] = transformed.arrow_table[col] + + updated_types = base.types.copy() + + for col in removed_cols: + if col in updated_types: + del updated_types[col] + + for col in replacement_cols: + if col in transformed.types: + updated_types[col] = transformed.types[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed.types: + updated_types[unique_col] = transformed.types[col] + + modified_dataset = modify_table(base, updated_arrays, types=updated_types) + modified_dataset = modified_dataset.select_columns(new_columns_order) + + return modified_dataset + + +class ApplyConverterSchema(BaseSchema): + converter: schema_field( + component_field(parent="BaseConverter"), + placeholder={"component": "ColumnRemover", "params": {}}, + description=MultilingualString( + en="Converter to apply, together with its own configuration.", + es="Convertidor a aplicar, junto con su propia configuración.", + pt="Conversor a aplicar, junto com a sua própria configuração.", + de="Anzuwendender Konverter samt seiner eigenen Konfiguration.", + zh="要应用的转换器及其自身的配置。", + ), + alias=MultilingualString( + en="Converter", + es="Convertidor", + pt="Conversor", + de="Konverter", + zh="转换器", + ), + ) # type: ignore + scope: schema_field( + none_type(dict), + placeholder={"columns": [], "rows": []}, + description=MultilingualString( + en="Part of the dataset the converter applies to: 'columns' is a " + "list of {'idx': n} and 'rows' a list of n, both 1-based. An " + "empty column list means every column.", + es="Parte del conjunto de datos a la que se aplica el convertidor: " + "'columns' es una lista de {'idx': n} y 'rows' una lista de n, " + "ambas con base 1. Una lista de columnas vacía significa todas " + "las columnas.", + pt="Parte do conjunto de dados à qual o conversor se aplica: " + "'columns' é uma lista de {'idx': n} e 'rows' uma lista de n, " + "ambas com base 1. Uma lista de colunas vazia significa todas " + "as colunas.", + de="Teil des Datensatzes, auf den der Konverter angewendet wird: " + "'columns' ist eine Liste von {'idx': n} und 'rows' eine Liste " + "von n, beide 1-basiert. Eine leere Spaltenliste bedeutet alle " + "Spalten.", + zh="转换器作用的数据集范围：'columns' 是 {'idx': n} 的列表，" + "'rows' 是 n 的列表，均从 1 开始。空的列列表表示所有列。", + ), + alias=MultilingualString( + en="Scope", + es="Alcance", + pt="Escopo", + de="Geltungsbereich", + zh="范围", + ), + ) # type: ignore + target: schema_field( + none_type(dict), + placeholder=None, + description=MultilingualString( + en="Target column handed to the converter as y, as {'idx': n} with " + "n 1-based. Leave empty for unsupervised converters.", + es="Columna objetivo entregada al convertidor como y, como " + "{'idx': n} con n en base 1. Dejar vacío para convertidores no " + "supervisados.", + pt="Coluna alvo entregue ao conversor como y, como {'idx': n} com " + "n com base 1. Deixe vazio para conversores não supervisionados.", + de="Zielspalte, die dem Konverter als y übergeben wird, als " + "{'idx': n} mit 1-basiertem n. Für unüberwachte Konverter leer " + "lassen.", + zh="作为 y 传给转换器的目标列，格式为 {'idx': n}，n 从 1 开始。" + "无监督转换器留空。", + ), + alias=MultilingualString( + en="Target column", + es="Columna objetivo", + pt="Coluna alvo", + de="Zielspalte", + zh="目标列", + ), + ) # type: ignore + + +class ApplyConverterUnit(BaseUnit): + """Fit one converter on the scoped data and transform the dataset with it. + + Reads ``dataset`` and writes ``dataset``: the same key on both sides, so any + number of these can be chained in one context, each one seeing what the + previous one produced. + + That is also why nothing about column identity ever crosses the context + boundary. The scope is expressed as 1-based column and row indexes, and + those indexes are resolved against ``dataset.column_names`` read at the top + of ``execute`` — the dataset as it is *right now*. A converter that renames, + drops or adds columns changes what index 3 means, so a resolved column list + published to the context would be stale for the very next converter. + """ + + SCHEMA = ApplyConverterSchema + + REQUIRES = ("dataset",) + PROVIDES = ("dataset",) + + def __init__(self, **config) -> None: + super().__init__(**config) + # Memoized on the instance, never in the context: two units of this + # class can live in the same context and they are different converters. + self._converter_class = None + + @property + def _converter_name(self) -> str: + return self.config["converter"]["component"] + + def _scope(self) -> Dict[str, List]: + """The configured scope, tolerating an explicit ``None``. + + The API schema always writes the ``scope`` key but allows it to be null, + so a converter saved without a scope arrives here as ``None`` rather + than as a missing key. + """ + return self.config.get("scope") or EMPTY_SCOPE + + def _target_index(self) -> int: + """The 1-based target column index, or None when there is no target.""" + target = self.config.get("target") + if target is None: + return None + return target.get("idx") + + def _check_target_bounds(self, dataset: "DashAIDataset") -> None: + """Reject a target index that does not point at a column. + + Checked against the dataset in hand rather than once up front: when + several converters are chained the column count changes underneath, and + an unchecked index would surface as a bare IndexError. + """ + target_column_index = self._target_index() + if target_column_index is None: + return + + if int(target_column_index) < 1 or int(target_column_index) > len( + dataset.features + ): + raise JobError( + f"Target column index {target_column_index} is out of bounds" + ) + + def _resolve_converter(self): + """Look the converter class up in the registry, once per instance.""" + if self._converter_class is not None: + return self._converter_class + + from kink import di + + component_registry = di["component_registry"] + converter_name = self._converter_name + + try: + self._converter_class = component_registry[converter_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Error importing converter {converter_name}: {e}") from e + + return self._converter_class + + def _resolve_scope( + self, dataset: "DashAIDataset" + ) -> Tuple[List[str], List[int], Any]: + """Turn the 1-based scope into names and indexes for this dataset.""" + scope = self._scope() + column_names = dataset.column_names + + columns_scope = [column["idx"] - 1 for column in scope["columns"]] + scope_column_indexes = sorted(set(columns_scope)) + + if not scope_column_indexes: + scope_column_indexes = list(range(len(dataset.features))) + + scope_column_names = [column_names[index] for index in scope_column_indexes] + + rows_scope = [row - 1 for row in scope["rows"]] + scope_rows_indexes = sorted(set(rows_scope)) + + target_column_name = None + target_column_index = self._target_index() + if target_column_index is not None: + target_column_name = column_names[int(target_column_index) - 1] + + return scope_column_names, scope_rows_indexes, target_column_name + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + loaded_dataset: "DashAIDataset" = ctx.require("dataset") + converter_name = self._converter_name + + converter_constructor = self._resolve_converter() + converter_instance = converter_constructor( + **(self.config["converter"].get("params") or {}) + ) + + self._check_target_bounds(loaded_dataset) + ( + scope_column_names, + scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(loaded_dataset) + + log.info(f"Applying converter: {converter_name}") + + y_dataset_fit = None + y_full_transform = None + if target_column_name is not None: + y_dataset_fit = loaded_dataset.select_columns([target_column_name]) + if scope_rows_indexes: + y_dataset_fit = y_dataset_fit.select(scope_rows_indexes) + y_full_transform = loaded_dataset.select_columns([target_column_name]) + else: + y_full_transform = y_dataset_fit + + X_dataset_fit = loaded_dataset.select_columns(scope_column_names) + + if scope_rows_indexes: + X_dataset_fit = X_dataset_fit.select(scope_rows_indexes) + + try: + converter_instance = converter_instance.fit(X_dataset_fit, y_dataset_fit) + except ValueError as e: + log.error(f"Validation error in {converter_name}: {e}") + raise JobError(f"Validation error fitting {converter_name}: {e}") from e + except Exception as e: + log.exception(e) + raise JobError(f"Error fitting converter {converter_name}: {e}") from e + + if scope_rows_indexes: + X_full_transform = loaded_dataset.select_columns(scope_column_names) + else: + # Deliberately the *same object* as the one passed to fit, not an + # equal one: converters such as TypeCast key a cache off the + # identity of the dataset they were fitted on and skip recomputing + # when transform receives it again. + X_full_transform = X_dataset_fit + + try: + transformed_dataset = converter_instance.transform( + X_full_transform, y_full_transform + ) + except Exception as e: + log.exception(e) + raise JobError(f"Error transforming data with {converter_name}: {e}") from e + + if type(converter_instance).CHANGES_ROW_COUNT: + loaded_dataset = transformed_dataset + else: + loaded_dataset = rebuild_dataset_with_transformed_columns( + loaded_dataset, + transformed_dataset, + scope_column_names, + ) + + ctx.put("dataset", loaded_dataset) diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py index 45c14c936..6579c3bee 100644 --- a/DashAI/back/units/build_model_unit.py +++ b/DashAI/back/units/build_model_unit.py @@ -120,7 +120,11 @@ class BuildModelUnit(BaseUnit): SCHEMA = BuildModelSchema - REQUIRES = ("x", "y", "n_labels") + # run_id and task_name only appear in the ModelFactory call and in error + # messages, but they are declared all the same: a key read without being + # declared is invisible to any caller — and to any future DAG validator — + # that inspects REQUIRES instead of running the unit. + REQUIRES = ("x", "y", "n_labels", "run_id", "task_name") PROVIDES = ("model", "factory", "optimizable_parameters", "model_parameters") def __init__(self, **config) -> None: @@ -195,6 +199,8 @@ def execute(self, ctx: ExecutionContext) -> None: component_registry = di["component_registry"] parameters = self.model_parameters + run_id = ctx.require("run_id") + task_name = ctx.require("task_name") model_class = self._resolve_model_class() @@ -212,15 +218,14 @@ def execute(self, ctx: ExecutionContext) -> None: except Exception as e: log.exception(e) raise JobError( - "Unable to find metrics associated with" - f"Task {ctx.get('task_name')} in registry", + f"Unable to find metrics associated withTask {task_name} in registry", ) from e try: factory = ModelFactory( model_class, parameters, - ctx.get("run_id"), + run_id, ctx.require("x"), ctx.require("y"), train_metrics, @@ -232,7 +237,7 @@ def execute(self, ctx: ExecutionContext) -> None: except Exception as e: log.exception(e) raise JobError( - f"Unable to instantiate model using run {ctx.get('run_id')}", + f"Unable to instantiate model using run {run_id}", ) from e # The original tree is what the search unit rewrites with the best diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 8fc16f958..8f9c6bf14 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -92,18 +92,26 @@ class FitModelUnit(BaseUnit): "x", "y", "task", + "run_id", ) PROVIDES = ("model", "plot_paths") - def validate(self, ctx: ExecutionContext) -> None: - # ctx.require, not ctx.get: "optimizable_parameters" is one of this - # unit's REQUIRES, so its absence means BuildModelUnit hasn't run yet - # — a call-order mistake, not "there is nothing to optimize". Only an - # empty value (the key present, genuinely no optimizable parameters) - # skips the optimizer/goal-metric checks below, so no registry lookup - # is needed either. - if not ctx.require("optimizable_parameters"): - return + def __init__(self, **config) -> None: + super().__init__(**config) + self._optimizer = None + self._goal_metric = None + + def _resolve_search(self): + """Resolve the optimizer and the goal metric, memoized on this unit. + + Kept on the instance rather than in the context on purpose. These are + this unit's own state, not something it hands to another unit: two + ``FitModelUnit`` instances sharing a context — a DAG with two training + nodes — would otherwise overwrite each other's optimizer, and the + second one would silently run the first one's. + """ + if self._optimizer is not None: + return self._optimizer, self._goal_metric from kink import di @@ -132,8 +140,21 @@ def validate(self, ctx: ExecutionContext) -> None: f"Error instantiating optimizer {optimizer_name}, {e}", ) from e - ctx.put("goal_metric", goal_metric) - ctx.put("optimizer", optimizer) + self._goal_metric = goal_metric + self._optimizer = optimizer + return optimizer, goal_metric + + def validate(self, ctx: ExecutionContext) -> None: + # ctx.require, not ctx.get: "optimizable_parameters" is one of this + # unit's REQUIRES, so its absence means BuildModelUnit hasn't run yet + # — a call-order mistake, not "there is nothing to optimize". Only an + # empty value (the key present, genuinely no optimizable parameters) + # skips the optimizer/goal-metric checks below, so no registry lookup + # is needed either. + if not ctx.require("optimizable_parameters"): + return + + self._resolve_search() def execute(self, ctx: ExecutionContext) -> None: import os @@ -146,6 +167,7 @@ def execute(self, ctx: ExecutionContext) -> None: model = ctx.require("model") x = ctx.require("x") y = ctx.require("y") + run_id = ctx.require("run_id") optimizable_parameters = ctx.require("optimizable_parameters") plot_paths = [] @@ -153,12 +175,10 @@ def execute(self, ctx: ExecutionContext) -> None: if not optimizable_parameters: model.train(x["train"], y["train"], x["validation"], y["validation"]) else: - # __call__ always runs validate() immediately before execute(), - # so "optimizer"/"goal_metric" are already in ctx here. - optimizer = ctx.require("optimizer") - goal_metric = ctx.require("goal_metric") + # Memoized: validate() resolved these already, and resolving + # again here would be the same lookup. + optimizer, goal_metric = self._resolve_search() factory = ctx.require("factory") - run_id = ctx.get("run_id") optimizer.optimize( model, diff --git a/DashAI/back/units/load_dataset_unit.py b/DashAI/back/units/load_dataset_unit.py index dc3244645..cac25e15c 100644 --- a/DashAI/back/units/load_dataset_unit.py +++ b/DashAI/back/units/load_dataset_unit.py @@ -3,9 +3,14 @@ import logging from typing import TYPE_CHECKING -from DashAI.back.core.schema_fields import BaseSchema, int_field, schema_field +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + none_type, + schema_field, +) from DashAI.back.core.utils import MultilingualString -from DashAI.back.dependencies.database.models import Dataset +from DashAI.back.dependencies.database.models import Dataset, Notebook from DashAI.back.job.base_job import JobError from DashAI.back.units.base_unit import BaseUnit from DashAI.back.units.context import ExecutionContext @@ -18,14 +23,18 @@ class LoadDatasetSchema(BaseSchema): dataset_id: schema_field( - int_field(gt=0), - placeholder=1, + none_type(int_field(gt=0)), + placeholder=None, description=MultilingualString( - en="Identifier of the stored dataset to load.", - es="Identificador del conjunto de datos almacenado a cargar.", - pt="Identificador do conjunto de dados armazenado a carregar.", - de="Kennung des zu ladenden gespeicherten Datensatzes.", - zh="要加载的已存储数据集的标识符。", + en="Identifier of the stored dataset to load. Mutually exclusive " + "with the notebook identifier.", + es="Identificador del conjunto de datos almacenado a cargar. " + "Excluyente con el identificador del cuaderno.", + pt="Identificador do conjunto de dados armazenado a carregar. " + "Mutuamente exclusivo com o identificador do caderno.", + de="Kennung des zu ladenden gespeicherten Datensatzes. Schließt " + "die Notebook-Kennung aus.", + zh="要加载的已存储数据集的标识符。与笔记本标识符互斥。", ), alias=MultilingualString( en="Dataset", @@ -35,18 +44,51 @@ class LoadDatasetSchema(BaseSchema): zh="数据集", ), ) # type: ignore + notebook_id: schema_field( + none_type(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="Identifier of the notebook whose working copy of the dataset " + "should be loaded. Mutually exclusive with the dataset identifier.", + es="Identificador del cuaderno cuya copia de trabajo del conjunto " + "de datos se debe cargar. Excluyente con el identificador del " + "conjunto de datos.", + pt="Identificador do caderno cuja cópia de trabalho do conjunto de " + "dados deve ser carregada. Mutuamente exclusivo com o " + "identificador do conjunto de dados.", + de="Kennung des Notebooks, dessen Arbeitskopie des Datensatzes " + "geladen werden soll. Schließt die Datensatz-Kennung aus.", + zh="要加载其数据集工作副本的笔记本标识符。与数据集标识符互斥。", + ), + alias=MultilingualString( + en="Notebook", + es="Cuaderno", + pt="Caderno", + de="Notebook", + zh="笔记本", + ), + ) # type: ignore class LoadDatasetUnit(BaseUnit): """Load a dataset from disk into the execution context. - Resolves the dataset row to find where it is stored and materialises it, so - downstream units receive a dataset instead of an identifier. + Takes exactly one of two starting points and materialises the dataset each + one points at, so downstream units receive a dataset instead of an + identifier: + + * ``dataset_id``: the stored dataset itself, read from ``Dataset.file_path``. + * ``notebook_id``: the notebook's own working copy, read from + ``Notebook.file_path``. A notebook holds a private copy precisely so + converters can rewrite it without touching the source dataset. + + Either way the unit publishes ``dataset_path``, which is where a later unit + has to write the dataset back for the change to be visible. """ SCHEMA = LoadDatasetSchema - PROVIDES = ("dataset",) + PROVIDES = ("dataset", "dataset_id", "dataset_path") def execute(self, ctx: ExecutionContext) -> None: from kink import di @@ -54,21 +96,43 @@ def execute(self, ctx: ExecutionContext) -> None: from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset session_factory = di["session_factory"] - dataset_id: int = self.config["dataset_id"] + + dataset_id = self.config.get("dataset_id") + notebook_id = self.config.get("notebook_id") + + if (dataset_id is None) == (notebook_id is None): + raise JobError( + "LoadDatasetUnit needs exactly one of dataset_id or notebook_id." + ) with session_factory() as db: - dataset: Dataset = db.get(Dataset, dataset_id) - if not dataset: - raise JobError(f"Dataset {dataset_id} does not exist in DB.") - file_path = dataset.file_path + if dataset_id is not None: + dataset: Dataset = db.get(Dataset, dataset_id) + if not dataset: + raise JobError(f"Dataset {dataset_id} does not exist in DB.") + file_path = dataset.file_path + else: + notebook: Notebook = db.get(Notebook, notebook_id) + if not notebook: + raise JobError(f"Notebook {notebook_id} does not exist in DB.") + file_path = notebook.file_path + # The notebook's copy still belongs to a source dataset, and + # downstream error messages identify the work by that id. + dataset_id = notebook.dataset_id + + dataset_path = f"{file_path}/dataset" try: - loaded_dataset: "DashAIDataset" = load_dataset(f"{file_path}/dataset") + loaded_dataset: "DashAIDataset" = load_dataset(dataset_path) except Exception as e: log.exception(e) raise JobError( f"Can not load dataset from path {file_path}", ) from e - ctx.put_ref("dataset_id", dataset_id) + if not loaded_dataset: + raise JobError(f"Dataset with path {dataset_path} not found") + ctx.put("dataset", loaded_dataset) + ctx.put_ref("dataset_id", dataset_id) + ctx.put_ref("dataset_path", dataset_path) diff --git a/DashAI/back/units/prepare_and_split_unit.py b/DashAI/back/units/prepare_and_split_unit.py index adfaf5928..17df6c52c 100644 --- a/DashAI/back/units/prepare_and_split_unit.py +++ b/DashAI/back/units/prepare_and_split_unit.py @@ -109,8 +109,11 @@ class PrepareAndSplitUnit(BaseUnit): SCHEMA = PrepareAndSplitSchema - REQUIRES = ("dataset",) - PROVIDES = ("x", "y", "n_labels", "task", "split_indexes") + # dataset_id is declared even though it only decorates an error message: + # an undeclared read is a contract a DAG validator cannot see, and the + # loaders publish it precisely so downstream units can name their input. + REQUIRES = ("dataset", "dataset_id") + PROVIDES = ("x", "y", "n_labels", "task", "split_indexes", "task_name") def execute(self, ctx: ExecutionContext) -> None: from kink import di @@ -129,6 +132,7 @@ def execute(self, ctx: ExecutionContext) -> None: splits = self.config["splits"] loaded_dataset = ctx.require("dataset") + dataset_id = ctx.require("dataset_id") try: task: "BaseTask" = component_registry[task_name]["class"]() @@ -170,7 +174,7 @@ def execute(self, ctx: ExecutionContext) -> None: except Exception as e: log.exception(e) raise JobError( - f"Can not prepare Dataset {ctx.get('dataset_id')} for Task {task_name}", + f"Can not prepare Dataset {dataset_id} for Task {task_name}", ) from e ctx.put_ref("task_name", task_name) diff --git a/DashAI/back/units/save_dataset_unit.py b/DashAI/back/units/save_dataset_unit.py new file mode 100644 index 000000000..50d9c9bc7 --- /dev/null +++ b/DashAI/back/units/save_dataset_unit.py @@ -0,0 +1,34 @@ +"""Unit that persists the dataset in the context back to disk.""" + +import logging + +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 SaveDatasetUnit(BaseUnit): + """Write the dataset back to the path it was loaded from. + + Takes no configuration: the destination is ``dataset_path``, published by + whichever unit loaded the dataset, so a save can never land somewhere the + load did not come from. Declares no outputs — its result is on disk, not in + the context. + """ + + REQUIRES = ("dataset", "dataset_path") + PROVIDES = () + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset + + dataset = ctx.require("dataset") + dataset_path = ctx.require("dataset_path") + + try: + save_dataset(dataset, dataset_path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not save dataset to path {dataset_path}") from e diff --git a/tests/back/api/test_converter_job.py b/tests/back/api/test_converter_job.py new file mode 100644 index 000000000..10fe1cdf1 --- /dev/null +++ b/tests/back/api/test_converter_job.py @@ -0,0 +1,314 @@ +"""End-to-end regression net for ``ConverterJob``. + +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 +column names 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. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ConverterStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import Converter +from DashAI.back.job.base_job import JobError +from DashAI.back.job.converter_job import ConverterJob + +IRIS_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", +] +IRIS_ROWS = 150 + + +@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 every converter run + mutates the notebook's copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "converter job test"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _create_converter(client, notebook_id, converter_name, scope=None, target=None): + """Create a Converter row through the API and return its id.""" + response = client.post( + "/api/v1/converter/", + json={ + "notebook_id": notebook_id, + "converter": converter_name, + "parameters": { + "order": 0, + "params": {}, + "scope": scope if scope is not None else {"columns": [], "rows": []}, + "target": target, + }, + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def _stored_converter(client, converter_id): + """Read the Converter row straight from the database.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + converter = db.get(Converter, converter_id) + return { + "status": converter.status, + "start_time": converter.start_time, + "end_time": converter.end_time, + } + + +def _notebook_dataset(notebook): + return load_dataset(f"{notebook['file_path']}/dataset") + + +def test_the_notebook_starts_as_an_untouched_copy_of_the_dataset(notebook): + """Guards the fixture itself: the assertions below mean nothing if the + notebook copy does not start out as the full iris dataset.""" + dataset = _notebook_dataset(notebook) + + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + +def test_converter_job_removes_the_scoped_column_and_finishes(client, notebook): + """The happy path, end to end: status transitions and the dataset on disk. + + ``ColumnRemover`` deletes whatever is in scope, so a one-column scope is + directly observable in the saved dataset. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 2}], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + stored = _stored_converter(client, converter_id) + assert stored["status"] == ConverterStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == [ + "SepalLengthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + assert len(dataset) == IRIS_ROWS + + +def test_two_converters_chained_see_the_columns_the_previous_one_left(client, notebook): + """Column scope is resolved by index against the *current* dataset. + + Each converter is its own row and its own job invocation, and the second + one's ``idx`` must be read against the four columns the first one left + behind, not against the original five. This is the behaviour any atomic + decomposition has to keep: nothing may cache the original column list. + """ + first = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + ) + ConverterJob(converter_id=first).run() + + assert _notebook_dataset(notebook).column_names == [ + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + + # idx 1 now means SepalWidthCm, not SepalLengthCm. + second = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + ) + ConverterJob(converter_id=second).run() + + assert _stored_converter(client, second)["status"] == ConverterStatus.FINISHED + assert _notebook_dataset(notebook).column_names == [ + "PetalLengthCm", + "PetalWidthCm", + "Species", + ] + + +def test_a_changes_row_count_converter_replaces_the_whole_dataset(client, notebook): + """``CHANGES_ROW_COUNT`` takes the transform output as the new dataset. + + ``NanRemover`` returns only the scoped columns, so the columns outside the + scope are dropped — the documented consequence of replacing the dataset + instead of merging the transformed columns back in. + """ + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}, {"idx": 5}], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == ["SepalLengthCm", "Species"] + # iris has no missing values, so no row is dropped. + assert len(dataset) == IRIS_ROWS + + +def test_an_empty_column_scope_means_every_column(client, notebook): + """An empty scope is not "no columns"; it is "all of them".""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [], "rows": []}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_a_row_scope_fits_on_the_subset_and_transforms_the_whole_dataset( + client, notebook +): + """A row scope narrows ``fit`` only; ``transform`` still sees every row.""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}], "rows": [1, 2, 3]}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + + dataset = _notebook_dataset(notebook) + assert dataset.column_names == ["SepalLengthCm"] + assert len(dataset) == IRIS_ROWS + + +def test_a_target_column_is_resolved_and_passed_to_the_converter(client, notebook): + """The supervised path builds y and hands it to fit/transform.""" + converter_id = _create_converter( + client, + notebook["id"], + "NanRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 5}, + ) + + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ( + ConverterStatus.FINISHED + ) + assert _notebook_dataset(notebook).column_names == ["SepalLengthCm"] + + +def test_a_missing_converter_row_reports_it_by_id(client, notebook): + with pytest.raises(JobError, match="Converter with id 999999 not found"): + ConverterJob(converter_id=999999).run() + + +def test_an_unknown_converter_name_fails_with_the_wrapped_message(client, notebook): + """The registry lookup error is reported inside the outer wrapper. + + Both halves matter: the import error names the culprit, the wrapper is what + the jobs UI actually shows. + """ + converter_id = _create_converter( + client, notebook["id"], "ThisConverterDoesNotExist" + ) + + with pytest.raises(JobError) as excinfo: + ConverterJob(converter_id=converter_id).run() + + message = str(excinfo.value) + assert "Error applying converters to dataset" in message + assert "Error importing converter ThisConverterDoesNotExist" in message + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + # The dataset must be left untouched when the converter never ran. + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_an_out_of_bounds_target_index_reports_cannot_load_dataset(client, notebook): + """The "out of bounds" text is swallowed; only the wrapper reaches the user. + + The inner JobError is re-caught by the surrounding ``except Exception`` and + replaced, surviving only as ``__cause__``. Locking this in because it is an + easy detail to "fix" by accident while refactoring. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 99}, + ) + + with pytest.raises(JobError, match="Cannot load dataset from") as excinfo: + ConverterJob(converter_id=converter_id).run() + + assert "Target column index 99 is out of bounds" in str(excinfo.value.__cause__) + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS + + +def test_a_failing_converter_leaves_the_dataset_untouched(client, notebook): + """``ColumnRemover`` raises when asked for a column that is not there. + + Reaching that requires a scope index past the end of the dataset, which is + exactly what a stale column list would produce in a chained run. + """ + converter_id = _create_converter( + client, + notebook["id"], + "ColumnRemover", + scope={"columns": [{"idx": 99}], "rows": []}, + ) + + with pytest.raises(JobError, match="Error applying converters to dataset"): + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index 87de3d44a..943662008 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -10,6 +10,8 @@ "FitModelUnit", "EvaluateModelUnit", "SaveModelUnit", + "ApplyConverterUnit", + "SaveDatasetUnit", } @@ -36,7 +38,10 @@ def test_units_expose_a_schema_the_front_can_render(units): def test_unit_schemas_describe_their_configuration(units): - assert "dataset_id" in units["LoadDatasetUnit"]["schema"]["properties"] + assert set(units["LoadDatasetUnit"]["schema"]["properties"]) == { + "dataset_id", + "notebook_id", + } assert set(units["PrepareAndSplitUnit"]["schema"]["properties"]) == { "task_name", "input_columns", @@ -45,6 +50,13 @@ def test_unit_schemas_describe_their_configuration(units): } assert "model" in units["BuildModelUnit"]["schema"]["properties"] assert "optimizer" in units["FitModelUnit"]["schema"]["properties"] + assert set(units["ApplyConverterUnit"]["schema"]["properties"]) == { + "converter", + "scope", + "target", + } + # SaveDatasetUnit is configuration-free: it saves where the load said. + assert units["SaveDatasetUnit"]["schema"]["properties"] == {} def test_component_fields_tell_the_front_which_components_to_offer(units): @@ -56,10 +68,13 @@ 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"] assert model["parent"] == "BaseModel" assert optimizer["parent"] == "BaseOptimizer" + assert converter["parent"] == "BaseConverter" assert set(model["properties"]) == {"component", "params"} + assert set(converter["properties"]) == {"component", "params"} def test_a_component_field_parent_resolves_to_real_components(client: TestClient): diff --git a/tests/back/units/test_apply_converter_unit.py b/tests/back/units/test_apply_converter_unit.py new file mode 100644 index 000000000..cc960504a --- /dev/null +++ b/tests/back/units/test_apply_converter_unit.py @@ -0,0 +1,388 @@ +"""Contract tests for ApplyConverterUnit, isolated from ConverterJob. + +The job only ever runs one of these per invocation, so an end-to-end run cannot +show whether the unit is safe to chain. These tests build the context by hand +and run several units against it, which is what the future DAG will do. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Integer +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +class _DropScopedColumns: + """Converter that removes whatever is in scope, like ColumnRemover.""" + + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + self.params = params + self.columns = [] + + def fit(self, x, y=None): + self.columns = x.column_names + return self + + def transform(self, x, y=None): + return x.remove_columns(self.columns) + + +class _ReplaceWithScope: + """Converter that replaces the dataset with the scoped columns only.""" + + CHANGES_ROW_COUNT = True + + def __init__(self, **params): + self.params = params + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + return x + + +class _RecordingConverter: + """Converter that records what fit and transform were handed.""" + + CHANGES_ROW_COUNT = True + + def __init__(self, **params): + self.params = params + self.fit_x = None + self.fit_y = None + self.transform_x = None + self.transform_y = None + + def fit(self, x, y=None): + self.fit_x = x + self.fit_y = y + return self + + def transform(self, x, y=None): + self.transform_x = x + self.transform_y = y + return x + + +class _FailingFit: + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + pass + + def fit(self, x, y=None): + raise ValueError("bad input") + + def transform(self, x, y=None): # pragma: no cover - never reached + return x + + +class _FailingTransform: + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + pass + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + raise RuntimeError("boom") + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "DropScopedColumns": {"class": _DropScopedColumns}, + "ReplaceWithScope": {"class": _ReplaceWithScope}, + "RecordingConverter": {"class": _RecordingConverter}, + "FailingFit": {"class": _FailingFit}, + "FailingTransform": {"class": _FailingTransform}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _unit(name, scope=None, target=None, params=None): + return ApplyConverterUnit( + converter={"component": name, "params": params or {}}, + scope=scope, + target=target, + ) + + +def _ctx(dataset): + ctx = ExecutionContext() + ctx.put("dataset", dataset) + return ctx + + +def test_the_unit_refuses_to_run_without_a_dataset(registry): + with pytest.raises(UnitContractError, match="'dataset' is not available"): + _unit("DropScopedColumns")(ExecutionContext()) + + +def test_a_scoped_column_is_resolved_by_one_based_index(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 2}], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["a", "c"] + + +def test_an_empty_column_scope_selects_every_column(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("ReplaceWithScope", scope={"columns": [], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_a_null_scope_is_treated_as_an_empty_one(registry): + """The API schema writes the ``scope`` key but allows it to be null. + + ``dict.get("scope", default)`` returns ``None`` rather than the default when + the key is present and null, so the unit has to coalesce it explicitly. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("ReplaceWithScope", scope=None)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_chained_units_resolve_their_indexes_against_the_current_dataset(registry): + """The reason no column identity may cross the context boundary. + + The first unit drops column ``a``, so index 1 means ``b`` by the time the + second unit runs. A column list resolved once and published to the context + would make the second unit drop the wrong column. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + assert ctx.require("dataset").column_names == ["b", "c"] + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + assert ctx.require("dataset").column_names == ["c"] + + +def test_the_unit_publishes_no_column_state_into_the_context(registry): + """Nothing resolved from the dataset may outlive the unit that resolved it.""" + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + + _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + for leaked in ( + "column_names", + "scope_column_names", + "scope_column_indexes", + "target_column_name", + "converter", + ): + assert not ctx.has(leaked), leaked + + +def test_two_units_in_one_context_do_not_share_their_resolved_converter(registry): + """Registry lookups are memoized on the instance, never in the context. + + A class cached under a context key would make the second unit silently run + the first one's converter. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + first = _unit("DropScopedColumns", scope={"columns": [{"idx": 1}], "rows": []}) + second = _unit("ReplaceWithScope", scope={"columns": [{"idx": 1}], "rows": []}) + + first(ctx) + second(ctx) + + assert first._converter_class is _DropScopedColumns + assert second._converter_class is _ReplaceWithScope + # ReplaceWithScope keeps only what is in scope, which is now "b". + assert ctx.require("dataset").column_names == ["b"] + + +def test_a_row_scope_narrows_fit_but_not_transform(registry): + ctx = _ctx(_dataset(a=[1, 2, 3, 4], b=[5, 6, 7, 8])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_rows"] = len(x) + return super().fit(x, y) + + def transform(self, x, y=None): + recorded["transform_rows"] = len(x) + recorded["same_object"] = x is self.fit_x + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": [1, 2]})(ctx) + + assert recorded["fit_rows"] == 2 + assert recorded["transform_rows"] == 4 + assert recorded["same_object"] is False + + +def test_no_row_scope_hands_transform_the_same_object_as_fit(registry): + """Object identity, not equality, is part of the contract here. + + ``TypeCastConverter`` caches converted columns during ``fit`` and reuses + them in ``transform`` only when handed the same dataset object. Losing that + identity does not fail anything — it just silently recomputes. + """ + ctx = _ctx(_dataset(a=[1, 2, 3, 4], b=[5, 6, 7, 8])) + recorded = {} + + class _Spy(_RecordingConverter): + def transform(self, x, y=None): + recorded["same_object"] = x is self.fit_x + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + assert recorded["same_object"] is True + + +def test_a_target_column_is_resolved_and_handed_over_as_y(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], label=[5, 6])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_y"] = None if y is None else y.column_names + return super().fit(x, y) + + def transform(self, x, y=None): + recorded["transform_y"] = None if y is None else y.column_names + return super().transform(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit( + "Spy", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 3}, + )(ctx) + + assert recorded["fit_y"] == ["label"] + assert recorded["transform_y"] == ["label"] + + +def test_no_target_means_no_y(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + recorded = {} + + class _Spy(_RecordingConverter): + def fit(self, x, y=None): + recorded["fit_y"] = y + return super().fit(x, y) + + registry["Spy"] = {"class": _Spy} + + _unit("Spy", scope={"columns": [{"idx": 1}], "rows": []}, target=None)(ctx) + + assert recorded["fit_y"] is None + + +def test_validate_rejects_an_out_of_bounds_target_without_running(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4])) + unit = _unit("DropScopedColumns", target={"idx": 99}) + + with pytest.raises(JobError, match="Target column index 99 is out of bounds"): + unit.validate(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + + +def test_the_target_bound_is_rechecked_against_the_dataset_of_the_moment(registry): + """A target that was valid before the previous converter ran may not be now. + + ``validate`` runs against whatever dataset is in the context when it is + called; chaining means ``execute`` can face a narrower one. Re-checking + turns what would be a bare IndexError into the same JobError. + """ + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + unit = _unit( + "DropScopedColumns", + scope={"columns": [{"idx": 1}], "rows": []}, + target={"idx": 3}, + ) + + unit.validate(ctx) # 3 columns: idx 3 is fine right now. + + # Another converter narrows the dataset behind its back. + _unit("ReplaceWithScope", scope={"columns": [{"idx": 1}], "rows": []})(ctx) + + with pytest.raises(JobError, match="Target column index 3 is out of bounds"): + unit(ctx) + + +def test_an_unknown_converter_names_the_culprit(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Error importing converter NotRegistered"): + _unit("NotRegistered")(ctx) + + +def test_a_value_error_during_fit_is_reported_as_a_validation_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Validation error fitting FailingFit"): + _unit("FailingFit")(ctx) + + +def test_any_other_error_during_fit_is_reported_as_a_fit_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + class _Boom(_FailingFit): + def fit(self, x, y=None): + raise RuntimeError("nope") + + registry["Boom"] = {"class": _Boom} + + with pytest.raises(JobError, match="Error fitting converter Boom"): + _unit("Boom")(ctx) + + +def test_an_error_during_transform_is_reported_as_a_transform_error(registry): + ctx = _ctx(_dataset(a=[1, 2])) + + with pytest.raises(JobError, match="Error transforming data with FailingTransform"): + _unit("FailingTransform")(ctx) + + +def test_converter_params_reach_the_constructor(registry): + ctx = _ctx(_dataset(a=[1, 2])) + unit = _unit("ReplaceWithScope", params={"threshold": 3}) + + unit(ctx) + + assert unit._converter_class is _ReplaceWithScope + + +def test_a_changes_row_count_converter_replaces_the_whole_dataset(registry): + ctx = _ctx(_dataset(a=[1, 2], b=[3, 4], c=[5, 6])) + + _unit("ReplaceWithScope", scope={"columns": [{"idx": 2}], "rows": []})(ctx) + + assert ctx.require("dataset").column_names == ["b"] diff --git a/tests/back/units/test_build_model_unit.py b/tests/back/units/test_build_model_unit.py index 8a4c3f624..a14aee9ab 100644 --- a/tests/back/units/test_build_model_unit.py +++ b/tests/back/units/test_build_model_unit.py @@ -51,6 +51,8 @@ def test_two_build_model_units_in_one_context_resolve_independently(fake_registr ctx.put("x", {"train": None, "validation": None}) ctx.put("y", {"train": None, "validation": None}) ctx.put("n_labels", None) + ctx.put_ref("run_id", 1) + ctx.put_ref("task_name", "ATask") a = _build_unit("ModelA") a(ctx) diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py index b4a0bab33..3ce22551c 100644 --- a/tests/back/units/test_fit_model_unit.py +++ b/tests/back/units/test_fit_model_unit.py @@ -1,6 +1,7 @@ """Tests for FitModelUnit's validation, independent of an actual training run.""" import pytest +from kink import di from DashAI.back.units.context import ExecutionContext, UnitContractError from DashAI.back.units.fit_model_unit import FitModelUnit @@ -34,7 +35,40 @@ def test_validate_is_a_noop_when_there_are_genuinely_no_optimizable_parameters() # Should not raise, and should not need the optimizer/goal_metric to # resolve in the registry. - _unit(optimizer_name="DoesNotExist", goal_metric="DoesNotExist").validate(ctx) + unit = _unit(optimizer_name="DoesNotExist", goal_metric="DoesNotExist") + unit.validate(ctx) - assert not ctx.has("optimizer") - assert not ctx.has("goal_metric") + assert unit._optimizer is None + assert unit._goal_metric is None + + +def test_the_optimizer_is_kept_on_the_unit_not_in_the_shared_context(): + """Regression: the optimizer is this unit's own state, not an output. + + It used to be written to the context by ``validate`` and read back by + ``execute``, using the shared context as a scratchpad between one unit's + own two phases. Two FitModelUnits in one context would overwrite each + other's optimizer, and the second would silently run the first one's. + """ + + class _Optimizer: + def __init__(self, **params): + pass + + registry = { + "AnOptimizer": {"class": _Optimizer}, + "Accuracy": {"class": object, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + try: + ctx = ExecutionContext() + ctx.put("optimizable_parameters", ["lr"]) + + unit = _unit(optimizer_name="AnOptimizer") + unit.validate(ctx) + + assert isinstance(unit._optimizer, _Optimizer) + assert not ctx.has("optimizer") + assert not ctx.has("goal_metric") + finally: + del di["component_registry"] diff --git a/tests/back/units/test_load_dataset_unit.py b/tests/back/units/test_load_dataset_unit.py new file mode 100644 index 000000000..1c5bbc994 --- /dev/null +++ b/tests/back/units/test_load_dataset_unit.py @@ -0,0 +1,151 @@ +"""Contract tests for LoadDatasetUnit, 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 pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit + + +class _Row: + """Stand-in for a Dataset or Notebook ORM row.""" + + def __init__(self, file_path=None, dataset_id=None): + self.file_path = file_path + self.dataset_id = dataset_id + + +class _FakeSession: + """Session that answers ``get`` from a table -> {id: row} mapping.""" + + 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) + + +@pytest.fixture(name="stored_dataset") +def fixture_stored_dataset(tmp_path): + """A real two-column dataset written to ``/store/dataset``.""" + 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]}) + types = { + "a": Integer(arrow_type=pa.int64()), + "b": Integer(arrow_type=pa.int64()), + } + root = tmp_path / "store" + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(stored_dataset): + """Point both a Dataset row and a Notebook row at the stored dataset.""" + rows = { + "Dataset": {7: _Row(file_path=str(stored_dataset))}, + "Notebook": {3: _Row(file_path=str(stored_dataset), dataset_id=7)}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def test_loading_by_dataset_id_publishes_the_whole_contract(fake_db): + ctx = ExecutionContext() + + LoadDatasetUnit(dataset_id=7)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + assert ctx.require("dataset_id") == 7 + assert ctx.require("dataset_path").endswith("dataset") + + +def test_loading_by_notebook_id_resolves_the_source_dataset_id(fake_db): + """The notebook branch must still publish a dataset id. + + Downstream error messages identify the work by dataset id, so a notebook + load that left the key unset would report ``None`` instead of the dataset. + """ + ctx = ExecutionContext() + + LoadDatasetUnit(notebook_id=3)(ctx) + + assert ctx.require("dataset").column_names == ["a", "b"] + assert ctx.require("dataset_id") == 7 + + +def test_the_two_starting_points_are_mutually_exclusive(fake_db): + with pytest.raises(JobError, match="exactly one of dataset_id or notebook_id"): + LoadDatasetUnit(dataset_id=7, notebook_id=3)(ExecutionContext()) + + +def test_no_starting_point_at_all_is_rejected(fake_db): + with pytest.raises(JobError, match="exactly one of dataset_id or notebook_id"): + LoadDatasetUnit()(ExecutionContext()) + + +def test_a_missing_dataset_row_is_reported_by_id(fake_db): + with pytest.raises(JobError, match="Dataset 99 does not exist in DB."): + LoadDatasetUnit(dataset_id=99)(ExecutionContext()) + + +def test_a_missing_notebook_row_is_reported_by_id(fake_db): + with pytest.raises(JobError, match="Notebook 99 does not exist in DB."): + LoadDatasetUnit(notebook_id=99)(ExecutionContext()) + + +def test_an_unreadable_path_becomes_a_job_error(fake_db, tmp_path): + fake_db["Notebook"][4] = _Row(file_path=str(tmp_path / "nowhere"), dataset_id=7) + + with pytest.raises(JobError, match="Can not load dataset from path"): + LoadDatasetUnit(notebook_id=4)(ExecutionContext()) + + +def test_the_dataset_is_cached_live_not_copied(fake_db): + """The dataset must come back as the same object, not a copy. + + ``ctx.get`` deep-copies the refs half and returns the cache half by + reference; a dataset that came back copied would mean every unit downstream + transformed a different object than the one that gets saved. + """ + ctx = ExecutionContext() + + LoadDatasetUnit(dataset_id=7)(ctx) + + assert ctx.require("dataset") is ctx.require("dataset") diff --git a/tests/back/units/test_save_dataset_unit.py b/tests/back/units/test_save_dataset_unit.py new file mode 100644 index 000000000..a1e534c08 --- /dev/null +++ b/tests/back/units/test_save_dataset_unit.py @@ -0,0 +1,64 @@ +"""Contract tests for SaveDatasetUnit.""" + +import pandas as pd +import pyarrow as pa +import pytest + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Integer +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def test_the_dataset_is_written_where_the_path_says(tmp_path): + destination = str(tmp_path / "notebook" / "dataset") + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1, 2], b=[3, 4])) + ctx.put_ref("dataset_path", destination) + + SaveDatasetUnit()(ctx) + + assert load_dataset(destination).column_names == ["a", "b"] + + +def test_saving_without_a_dataset_is_a_contract_error(tmp_path): + ctx = ExecutionContext() + ctx.put_ref("dataset_path", str(tmp_path / "dataset")) + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + SaveDatasetUnit()(ctx) + + +def test_saving_without_a_path_is_a_contract_error(): + """A missing path is a wiring mistake, not a "nowhere to save" decision. + + The unit has no fallback destination on purpose: silently picking one would + write the dataset somewhere nobody asked for. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + + with pytest.raises(UnitContractError, match="'dataset_path' is not available"): + SaveDatasetUnit()(ctx) + + +def test_an_unwritable_destination_becomes_a_job_error(tmp_path): + blocker = tmp_path / "blocker" + blocker.write_text("not a directory", encoding="utf-8") + + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put_ref("dataset_path", str(blocker / "dataset")) + + with pytest.raises(JobError, match="Can not save dataset to path"): + SaveDatasetUnit()(ctx) diff --git a/tests/back/units/test_unit_contracts.py b/tests/back/units/test_unit_contracts.py new file mode 100644 index 000000000..b472a885a --- /dev/null +++ b/tests/back/units/test_unit_contracts.py @@ -0,0 +1,135 @@ +"""A contract audit over every registered unit, enforced as a test. + +The individual unit tests check behaviour; this one checks that the *declared* +contract matches the code. Undeclared context reads are the recurring mistake in +this design: they never break the job that happens to wire the context by hand, +so they survive every end-to-end test and only surface when something reuses the +unit — which is the whole point of having units. +""" + +import ast +import pathlib + +import pytest + +UNITS_DIR = pathlib.Path(__file__).resolve().parents[3] / "DashAI" / "back" / "units" + +#: Keys a unit reads that its own execution produces, so they need no declaration. +SELF_PRODUCED = {"dataset"} + + +def _unit_modules(): + for path in sorted(UNITS_DIR.glob("*.py")): + if path.name in {"__init__.py", "base_unit.py", "context.py"}: + continue + yield path + + +def _string_literals(node): + return { + element.value + for element in ast.walk(node) + if isinstance(element, ast.Constant) and isinstance(element.value, str) + } + + +def _unit_class(tree): + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and any( + isinstance(base, ast.Name) and base.id == "BaseUnit" for base in node.bases + ): + return node + return None + + +def _declared(cls, name): + for node in cls.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == name for t in node.targets + ): + return _string_literals(node.value) + return set() + + +def _context_calls(cls, methods): + """Every ``ctx.("key")`` literal inside the class.""" + keys = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in methods + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "ctx" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + keys.add(node.args[0].value) + return keys + + +def _parsed_units(): + units = [] + for path in _unit_modules(): + tree = ast.parse(path.read_text(encoding="utf-8")) + cls = _unit_class(tree) + if cls is not None: + units.append((path.name, cls)) + return units + + +UNITS = _parsed_units() + + +def test_the_audit_actually_found_the_units(): + """Guards the audit itself: a broken parser would make it vacuously pass.""" + assert {name for name, _ in UNITS} >= { + "load_dataset_unit.py", + "apply_converter_unit.py", + "save_dataset_unit.py", + } + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_every_context_key_a_unit_reads_is_declared_in_requires(name, cls): + read = _context_calls(cls, {"require", "get", "has"}) + declared = _declared(cls, "REQUIRES") | _declared(cls, "PROVIDES") | SELF_PRODUCED + + undeclared = read - declared + assert not undeclared, ( + f"{name} reads {sorted(undeclared)} from the context without declaring " + "them in REQUIRES. A caller inspecting the contract cannot know they " + "are needed, and a missing value reads as 'not applicable' instead of " + "'wiring mistake'." + ) + + +@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"}) + promised = _declared(cls, "PROVIDES") + + unwritten = promised - written + assert not unwritten, ( + f"{name} promises {sorted(unwritten)} in PROVIDES but never writes it." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_use_the_context_as_its_own_scratchpad(name, cls): + """A key written and read by the same unit, and promised to nobody. + + That is instance state wearing a context key's clothes: two units of the + same class in one context would overwrite each other. It belongs on + ``self``, memoized, the way the registry lookups are. + """ + written = _context_calls(cls, {"put", "put_ref"}) + read = _context_calls(cls, {"require", "get", "has"}) + promised = _declared(cls, "PROVIDES") + + scratch = (written & read) - promised - SELF_PRODUCED + assert not scratch, ( + f"{name} writes and reads {sorted(scratch)} without promising it. " + "Keep per-instance state on the unit, not in the shared context." + ) From f5b26c6b895e18422c7c9083600de7d32fbf9be8 Mon Sep 17 00:00:00 2001 From: Felipe Date: Sun, 2 Aug 2026 15:35:47 -0400 Subject: [PATCH 04/28] feat: Improve error handling in ConverterJob and add test for dataset load failure --- DashAI/back/job/converter_job.py | 10 ++++++++++ DashAI/back/units/apply_converter_unit.py | 6 +++--- DashAI/back/units/build_model_unit.py | 2 +- tests/back/api/test_converter_job.py | 22 ++++++++++++++++++++++ 4 files changed, 36 insertions(+), 4 deletions(-) diff --git a/DashAI/back/job/converter_job.py b/DashAI/back/job/converter_job.py index f4b8e187b..f833dd3b7 100644 --- a/DashAI/back/job/converter_job.py +++ b/DashAI/back/job/converter_job.py @@ -142,6 +142,16 @@ def run( converter.set_status_as_error() db.commit() raise JobError("Error loading dataset info") from e + except Exception: + # Anything the load unit raises (missing notebook, unreadable + # dataset) also has to leave the row in ERROR. Nothing marks it + # otherwise: the Huey error signal only writes to its own + # task_copy table, never to the Converter row, so without this + # the converter would stay STARTED forever. Re-raised as-is so + # the unit's specific message survives. + converter.set_status_as_error() + db.commit() + raise apply_converter = ApplyConverterUnit( converter={ diff --git a/DashAI/back/units/apply_converter_unit.py b/DashAI/back/units/apply_converter_unit.py index ab4c75998..4c4a97aba 100644 --- a/DashAI/back/units/apply_converter_unit.py +++ b/DashAI/back/units/apply_converter_unit.py @@ -1,7 +1,7 @@ """Unit that applies a single converter to the dataset in the context.""" import logging -from typing import TYPE_CHECKING, Any, Dict, List, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple from DashAI.back.core.schema_fields import ( BaseSchema, @@ -231,7 +231,7 @@ def _scope(self) -> Dict[str, List]: """ return self.config.get("scope") or EMPTY_SCOPE - def _target_index(self) -> int: + def _target_index(self) -> Optional[int]: """The 1-based target column index, or None when there is no target.""" target = self.config.get("target") if target is None: @@ -276,7 +276,7 @@ def _resolve_converter(self): def _resolve_scope( self, dataset: "DashAIDataset" - ) -> Tuple[List[str], List[int], Any]: + ) -> Tuple[List[str], List[int], Optional[str]]: """Turn the 1-based scope into names and indexes for this dataset.""" scope = self._scope() column_names = dataset.column_names diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py index 6579c3bee..3f39bb1b6 100644 --- a/DashAI/back/units/build_model_unit.py +++ b/DashAI/back/units/build_model_unit.py @@ -218,7 +218,7 @@ def execute(self, ctx: ExecutionContext) -> None: except Exception as e: log.exception(e) raise JobError( - f"Unable to find metrics associated withTask {task_name} in registry", + f"Unable to find metrics associated with Task {task_name} in registry", ) from e try: diff --git a/tests/back/api/test_converter_job.py b/tests/back/api/test_converter_job.py index 10fe1cdf1..925113b65 100644 --- a/tests/back/api/test_converter_job.py +++ b/tests/back/api/test_converter_job.py @@ -294,6 +294,28 @@ def test_an_out_of_bounds_target_index_reports_cannot_load_dataset(client, noteb assert _notebook_dataset(notebook).column_names == IRIS_COLUMNS +def test_a_dataset_that_cannot_be_loaded_still_leaves_the_row_in_error( + client, notebook +): + """A load failure must not leave the converter stuck in STARTED. + + Nothing else would fix it: the Huey error signal writes only to its own + ``task_copy`` table and never touches the ``Converter`` row, and the job + runs with no outer handler. Before this, only ``SQLAlchemyError`` was + caught here, so an unreadable dataset left the row STARTED forever and the + UI showed the converter as still running. + """ + import shutil + + converter_id = _create_converter(client, notebook["id"], "ColumnRemover") + shutil.rmtree(f"{notebook['file_path']}/dataset") + + with pytest.raises(JobError, match="Can not load dataset from path"): + ConverterJob(converter_id=converter_id).run() + + assert _stored_converter(client, converter_id)["status"] == ConverterStatus.ERROR + + def test_a_failing_converter_leaves_the_dataset_untouched(client, notebook): """``ColumnRemover`` raises when asked for a column that is not there. From 39f556dbdbec0dab1f0794200bbd2eb011e060db Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 3 Aug 2026 19:59:18 -0400 Subject: [PATCH 05/28] feat: Split ApplyConverterUnit into FitConverterUnit and TransformDatasetUnit - Introduced FitConverterUnit to fit a converter on a dataset without transforming it. - Introduced TransformDatasetUnit to apply an already fitted converter to a dataset. - Updated ApplyConverterUnit to utilize the new units for fitting and transforming. - Added ConverterScopeMixin for shared scope resolution logic between converter units. - Updated initial_components.py to include new units. - Added tests to ensure correct functionality of the new units and their interactions. --- DashAI/back/initial_components.py | 4 + DashAI/back/units/apply_converter_unit.py | 331 +++------------- DashAI/back/units/converter_scope.py | 357 ++++++++++++++++++ DashAI/back/units/fit_converter_unit.py | 107 ++++++ DashAI/back/units/transform_dataset_unit.py | 82 ++++ tests/back/api/test_units_api.py | 12 + .../test_converter_fit_transform_split.py | 278 ++++++++++++++ 7 files changed, 891 insertions(+), 280 deletions(-) create mode 100644 DashAI/back/units/converter_scope.py create mode 100644 DashAI/back/units/fit_converter_unit.py create mode 100644 DashAI/back/units/transform_dataset_unit.py create mode 100644 tests/back/units/test_converter_fit_transform_split.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 5cbcb849a..db86e671f 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -350,11 +350,13 @@ from DashAI.back.units.apply_converter_unit import ApplyConverterUnit from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_converter_unit import FitConverterUnit from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.load_dataset_unit import LoadDatasetUnit from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit from DashAI.back.units.save_dataset_unit import SaveDatasetUnit from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -526,6 +528,8 @@ def get_initial_components(): EvaluateModelUnit, SaveModelUnit, ApplyConverterUnit, + FitConverterUnit, + TransformDatasetUnit, SaveDatasetUnit, # Explainers ContrastiveShap, diff --git a/DashAI/back/units/apply_converter_unit.py b/DashAI/back/units/apply_converter_unit.py index 4c4a97aba..da4bc636d 100644 --- a/DashAI/back/units/apply_converter_unit.py +++ b/DashAI/back/units/apply_converter_unit.py @@ -1,199 +1,43 @@ -"""Unit that applies a single converter to the dataset in the context.""" +"""Unit that fits a converter and transforms one dataset with it.""" import logging -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING -from DashAI.back.core.schema_fields import ( - BaseSchema, - component_field, - none_type, - schema_field, -) -from DashAI.back.core.utils import MultilingualString +from DashAI.back.core.schema_fields import BaseSchema 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.converter_scope import ( + ConverterScopeMixin, + converter_field, + scope_field, + target_field, +) if TYPE_CHECKING: from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset log = logging.getLogger(__name__) -EMPTY_SCOPE: Dict[str, List] = {"columns": [], "rows": []} - - -def rebuild_dataset_with_transformed_columns( - base: "DashAIDataset", - transformed: "DashAIDataset", - scope_column_names: List[str], -) -> "DashAIDataset": - """ - Replaces specific columns in the base dataset with columns from the transformed - dataset, preserving their original positions. Also appends any additional columns - that were generated by the transformer at the end. Keeps the features and metadata - consistent. - - Parameters - ---------- - base : DashAIDataset - The original dataset before transformation. - - transformed : DashAIDataset - The dataset resulting from applying a transformer, containing updated and/or - new columns. - - scope_column_names : List[str] - Names of the columns that were originally selected for transformation. - - Returns - ------- - DashAIDataset - A new dataset with the specified columns replaced in place, new columns - appended, and original metadata and split information preserved. - """ - from DashAI.back.dataloaders.classes.dashai_dataset import modify_table - - original_columns = base.column_names - transformed_cols = transformed.column_names - - transformed_cols_set = set(transformed_cols) - scope_column_names_set = set(scope_column_names) - - removed_cols = [ - col for col in scope_column_names if col not in transformed_cols_set - ] - replacement_cols = [ - col for col in scope_column_names if col in transformed_cols_set - ] - new_cols = [col for col in transformed_cols if col not in scope_column_names_set] - - removed_cols_set = set(removed_cols) - new_columns_order = [] - seen_cols = set() - for col in original_columns: - if col in removed_cols_set: - continue - if col not in seen_cols: - new_columns_order.append(col) - seen_cols.add(col) - - col_name_mapping = {} - for col in new_cols: - unique_col = col - counter = 1 - while unique_col in seen_cols: - unique_col = f"{col}_{counter}" - counter += 1 - new_columns_order.append(unique_col) - seen_cols.add(unique_col) - col_name_mapping[col] = unique_col - - updated_arrays = {} - for col in replacement_cols: - if col in transformed_cols_set: - updated_arrays[col] = transformed.arrow_table[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed_cols_set: - updated_arrays[unique_col] = transformed.arrow_table[col] - - updated_types = base.types.copy() - - for col in removed_cols: - if col in updated_types: - del updated_types[col] - - for col in replacement_cols: - if col in transformed.types: - updated_types[col] = transformed.types[col] - for col, unique_col in col_name_mapping.items(): - if col in transformed.types: - updated_types[unique_col] = transformed.types[col] - - modified_dataset = modify_table(base, updated_arrays, types=updated_types) - modified_dataset = modified_dataset.select_columns(new_columns_order) - - return modified_dataset +class ApplyConverterSchema(BaseSchema): + converter: converter_field() # type: ignore + scope: scope_field() # type: ignore + target: target_field() # type: ignore -class ApplyConverterSchema(BaseSchema): - converter: schema_field( - component_field(parent="BaseConverter"), - placeholder={"component": "ColumnRemover", "params": {}}, - description=MultilingualString( - en="Converter to apply, together with its own configuration.", - es="Convertidor a aplicar, junto con su propia configuración.", - pt="Conversor a aplicar, junto com a sua própria configuração.", - de="Anzuwendender Konverter samt seiner eigenen Konfiguration.", - zh="要应用的转换器及其自身的配置。", - ), - alias=MultilingualString( - en="Converter", - es="Convertidor", - pt="Conversor", - de="Konverter", - zh="转换器", - ), - ) # type: ignore - scope: schema_field( - none_type(dict), - placeholder={"columns": [], "rows": []}, - description=MultilingualString( - en="Part of the dataset the converter applies to: 'columns' is a " - "list of {'idx': n} and 'rows' a list of n, both 1-based. An " - "empty column list means every column.", - es="Parte del conjunto de datos a la que se aplica el convertidor: " - "'columns' es una lista de {'idx': n} y 'rows' una lista de n, " - "ambas con base 1. Una lista de columnas vacía significa todas " - "las columnas.", - pt="Parte do conjunto de dados à qual o conversor se aplica: " - "'columns' é uma lista de {'idx': n} e 'rows' uma lista de n, " - "ambas com base 1. Uma lista de colunas vazia significa todas " - "as colunas.", - de="Teil des Datensatzes, auf den der Konverter angewendet wird: " - "'columns' ist eine Liste von {'idx': n} und 'rows' eine Liste " - "von n, beide 1-basiert. Eine leere Spaltenliste bedeutet alle " - "Spalten.", - zh="转换器作用的数据集范围：'columns' 是 {'idx': n} 的列表，" - "'rows' 是 n 的列表，均从 1 开始。空的列列表表示所有列。", - ), - alias=MultilingualString( - en="Scope", - es="Alcance", - pt="Escopo", - de="Geltungsbereich", - zh="范围", - ), - ) # type: ignore - target: schema_field( - none_type(dict), - placeholder=None, - description=MultilingualString( - en="Target column handed to the converter as y, as {'idx': n} with " - "n 1-based. Leave empty for unsupervised converters.", - es="Columna objetivo entregada al convertidor como y, como " - "{'idx': n} con n en base 1. Dejar vacío para convertidores no " - "supervisados.", - pt="Coluna alvo entregue ao conversor como y, como {'idx': n} com " - "n com base 1. Deixe vazio para conversores não supervisionados.", - de="Zielspalte, die dem Konverter als y übergeben wird, als " - "{'idx': n} mit 1-basiertem n. Für unüberwachte Konverter leer " - "lassen.", - zh="作为 y 传给转换器的目标列，格式为 {'idx': n}，n 从 1 开始。" - "无监督转换器留空。", - ), - alias=MultilingualString( - en="Target column", - es="Columna objetivo", - pt="Coluna alvo", - de="Zielspalte", - zh="目标列", - ), - ) # type: ignore +class ApplyConverterUnit(BaseUnit, ConverterScopeMixin): + """Fit one converter on the scoped data and transform the same dataset. + The single-dataset case, which is what applying a converter to a notebook + means. Equivalent to :class:`FitConverterUnit` followed by + :class:`TransformDatasetUnit` on the same dataset, kept as one unit for two + reasons: it is the common case, and when there is no row scope it hands + ``transform`` the very object ``fit`` saw, which some converters rely on to + skip recomputing (see below). -class ApplyConverterUnit(BaseUnit): - """Fit one converter on the scoped data and transform the dataset with it. + Use the split pair instead whenever the converter has to be learned from one + dataset and applied to another — fit on train, transform on test. Reads ``dataset`` and writes ``dataset``: the same key on both sides, so any number of these can be chained in one context, each one seeing what the @@ -210,7 +54,7 @@ class ApplyConverterUnit(BaseUnit): SCHEMA = ApplyConverterSchema REQUIRES = ("dataset",) - PROVIDES = ("dataset",) + PROVIDES = ("dataset", "fitted_converter") def __init__(self, **config) -> None: super().__init__(**config) @@ -222,40 +66,6 @@ def __init__(self, **config) -> None: def _converter_name(self) -> str: return self.config["converter"]["component"] - def _scope(self) -> Dict[str, List]: - """The configured scope, tolerating an explicit ``None``. - - The API schema always writes the ``scope`` key but allows it to be null, - so a converter saved without a scope arrives here as ``None`` rather - than as a missing key. - """ - return self.config.get("scope") or EMPTY_SCOPE - - def _target_index(self) -> Optional[int]: - """The 1-based target column index, or None when there is no target.""" - target = self.config.get("target") - if target is None: - return None - return target.get("idx") - - def _check_target_bounds(self, dataset: "DashAIDataset") -> None: - """Reject a target index that does not point at a column. - - Checked against the dataset in hand rather than once up front: when - several converters are chained the column count changes underneath, and - an unchecked index would surface as a bare IndexError. - """ - target_column_index = self._target_index() - if target_column_index is None: - return - - if int(target_column_index) < 1 or int(target_column_index) > len( - dataset.features - ): - raise JobError( - f"Target column index {target_column_index} is out of bounds" - ) - def _resolve_converter(self): """Look the converter class up in the registry, once per instance.""" if self._converter_class is not None: @@ -274,31 +84,6 @@ def _resolve_converter(self): return self._converter_class - def _resolve_scope( - self, dataset: "DashAIDataset" - ) -> Tuple[List[str], List[int], Optional[str]]: - """Turn the 1-based scope into names and indexes for this dataset.""" - scope = self._scope() - column_names = dataset.column_names - - columns_scope = [column["idx"] - 1 for column in scope["columns"]] - scope_column_indexes = sorted(set(columns_scope)) - - if not scope_column_indexes: - scope_column_indexes = list(range(len(dataset.features))) - - scope_column_names = [column_names[index] for index in scope_column_indexes] - - rows_scope = [row - 1 for row in scope["rows"]] - scope_rows_indexes = sorted(set(rows_scope)) - - target_column_name = None - target_column_index = self._target_index() - if target_column_index is not None: - target_column_name = column_names[int(target_column_index) - 1] - - return scope_column_names, scope_rows_indexes, target_column_name - def validate(self, ctx: ExecutionContext) -> None: """Reject an out-of-bounds target before any work is done.""" self._check_target_bounds(ctx.require("dataset")) @@ -321,54 +106,40 @@ def execute(self, ctx: ExecutionContext) -> None: log.info(f"Applying converter: {converter_name}") - y_dataset_fit = None - y_full_transform = None - if target_column_name is not None: - y_dataset_fit = loaded_dataset.select_columns([target_column_name]) - if scope_rows_indexes: - y_dataset_fit = y_dataset_fit.select(scope_rows_indexes) - y_full_transform = loaded_dataset.select_columns([target_column_name]) - else: - y_full_transform = y_dataset_fit - - X_dataset_fit = loaded_dataset.select_columns(scope_column_names) - - if scope_rows_indexes: - X_dataset_fit = X_dataset_fit.select(scope_rows_indexes) + x_dataset_fit, y_dataset_fit = self._slice_for_fit( + loaded_dataset, + scope_column_names, + scope_rows_indexes, + target_column_name, + ) - try: - converter_instance = converter_instance.fit(X_dataset_fit, y_dataset_fit) - except ValueError as e: - log.error(f"Validation error in {converter_name}: {e}") - raise JobError(f"Validation error fitting {converter_name}: {e}") from e - except Exception as e: - log.exception(e) - raise JobError(f"Error fitting converter {converter_name}: {e}") from e + converter_instance = self._fit( + converter_instance, x_dataset_fit, y_dataset_fit, converter_name + ) if scope_rows_indexes: - X_full_transform = loaded_dataset.select_columns(scope_column_names) + x_full_transform, y_full_transform = self._slice_for_transform( + loaded_dataset, scope_column_names, target_column_name + ) else: - # Deliberately the *same object* as the one passed to fit, not an - # equal one: converters such as TypeCast key a cache off the + # Deliberately the *same objects* as the ones passed to fit, not + # equal ones: converters such as TypeCast key a cache off the # identity of the dataset they were fitted on and skip recomputing - # when transform receives it again. - X_full_transform = X_dataset_fit + # when transform receives it again. Without a row scope the fit + # slice already covers every row, so reusing it is also correct. + x_full_transform, y_full_transform = x_dataset_fit, y_dataset_fit - try: - transformed_dataset = converter_instance.transform( - X_full_transform, y_full_transform - ) - except Exception as e: - log.exception(e) - raise JobError(f"Error transforming data with {converter_name}: {e}") from e + transformed_dataset = self._transform( + converter_instance, x_full_transform, y_full_transform, converter_name + ) - if type(converter_instance).CHANGES_ROW_COUNT: - loaded_dataset = transformed_dataset - else: - loaded_dataset = rebuild_dataset_with_transformed_columns( + ctx.put( + "dataset", + self._merge_transformed( + converter_instance, loaded_dataset, transformed_dataset, scope_column_names, - ) - - ctx.put("dataset", loaded_dataset) + ), + ) + ctx.put("fitted_converter", converter_instance) diff --git a/DashAI/back/units/converter_scope.py b/DashAI/back/units/converter_scope.py new file mode 100644 index 000000000..73a10a712 --- /dev/null +++ b/DashAI/back/units/converter_scope.py @@ -0,0 +1,357 @@ +"""Shared scope resolution and slicing for the converter units. + +The three converter units — fit, transform, and the fused apply — all express +which part of a dataset they touch the same way, and all resolve it the same +way. That logic lives here so the fused path and the split path cannot drift +apart. +""" + +import logging +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +from DashAI.back.core.schema_fields import ( + component_field, + none_type, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + +EMPTY_SCOPE: Dict[str, List] = {"columns": [], "rows": []} + + +def converter_field(): + """Schema field for picking a converter and configuring it.""" + return schema_field( + component_field(parent="BaseConverter"), + placeholder={"component": "ColumnRemover", "params": {}}, + description=MultilingualString( + en="Converter to apply, together with its own configuration.", + es="Convertidor a aplicar, junto con su propia configuración.", + pt="Conversor a aplicar, junto com a sua própria configuração.", + de="Anzuwendender Konverter samt seiner eigenen Konfiguration.", + zh="要应用的转换器及其自身的配置。", + ), + alias=MultilingualString( + en="Converter", + es="Convertidor", + pt="Conversor", + de="Konverter", + zh="转换器", + ), + ) + + +def scope_field(): + """Schema field for the part of the dataset a converter touches.""" + return schema_field( + none_type(dict), + placeholder={"columns": [], "rows": []}, + description=MultilingualString( + en="Part of the dataset the converter applies to: 'columns' is a " + "list of {'idx': n} and 'rows' a list of n, both 1-based. An " + "empty column list means every column.", + es="Parte del conjunto de datos a la que se aplica el convertidor: " + "'columns' es una lista de {'idx': n} y 'rows' una lista de n, " + "ambas con base 1. Una lista de columnas vacía significa todas " + "las columnas.", + pt="Parte do conjunto de dados à qual o conversor se aplica: " + "'columns' é uma lista de {'idx': n} e 'rows' uma lista de n, " + "ambas com base 1. Uma lista de colunas vazia significa todas " + "as colunas.", + de="Teil des Datensatzes, auf den der Konverter angewendet wird: " + "'columns' ist eine Liste von {'idx': n} und 'rows' eine Liste " + "von n, beide 1-basiert. Eine leere Spaltenliste bedeutet alle " + "Spalten.", + zh="转换器作用的数据集范围：'columns' 是 {'idx': n} 的列表，" + "'rows' 是 n 的列表，均从 1 开始。空的列列表表示所有列。", + ), + alias=MultilingualString( + en="Scope", + es="Alcance", + pt="Escopo", + de="Geltungsbereich", + zh="范围", + ), + ) + + +def target_field(): + """Schema field for the column handed to the converter as ``y``.""" + return schema_field( + none_type(dict), + placeholder=None, + description=MultilingualString( + en="Target column handed to the converter as y, as {'idx': n} with " + "n 1-based. Leave empty for unsupervised converters.", + es="Columna objetivo entregada al convertidor como y, como " + "{'idx': n} con n en base 1. Dejar vacío para convertidores no " + "supervisados.", + pt="Coluna alvo entregue ao conversor como y, como {'idx': n} com " + "n com base 1. Deixe vazio para conversores não supervisionados.", + de="Zielspalte, die dem Konverter als y übergeben wird, als " + "{'idx': n} mit 1-basiertem n. Für unüberwachte Konverter leer " + "lassen.", + zh="作为 y 传给转换器的目标列，格式为 {'idx': n}，n 从 1 开始。" + "无监督转换器留空。", + ), + alias=MultilingualString( + en="Target column", + es="Columna objetivo", + pt="Coluna alvo", + de="Zielspalte", + zh="目标列", + ), + ) + + +def rebuild_dataset_with_transformed_columns( + base: "DashAIDataset", + transformed: "DashAIDataset", + scope_column_names: List[str], +) -> "DashAIDataset": + """ + Replaces specific columns in the base dataset with columns from the transformed + dataset, preserving their original positions. Also appends any additional columns + that were generated by the transformer at the end. Keeps the features and metadata + consistent. + + Parameters + ---------- + base : DashAIDataset + The original dataset before transformation. + + transformed : DashAIDataset + The dataset resulting from applying a transformer, containing updated and/or + new columns. + + scope_column_names : List[str] + Names of the columns that were originally selected for transformation. + + Returns + ------- + DashAIDataset + A new dataset with the specified columns replaced in place, new columns + appended, and original metadata and split information preserved. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import modify_table + + original_columns = base.column_names + transformed_cols = transformed.column_names + + transformed_cols_set = set(transformed_cols) + scope_column_names_set = set(scope_column_names) + + removed_cols = [ + col for col in scope_column_names if col not in transformed_cols_set + ] + replacement_cols = [ + col for col in scope_column_names if col in transformed_cols_set + ] + new_cols = [col for col in transformed_cols if col not in scope_column_names_set] + + removed_cols_set = set(removed_cols) + + new_columns_order = [] + seen_cols = set() + for col in original_columns: + if col in removed_cols_set: + continue + if col not in seen_cols: + new_columns_order.append(col) + seen_cols.add(col) + + col_name_mapping = {} + for col in new_cols: + unique_col = col + counter = 1 + while unique_col in seen_cols: + unique_col = f"{col}_{counter}" + counter += 1 + new_columns_order.append(unique_col) + seen_cols.add(unique_col) + col_name_mapping[col] = unique_col + + updated_arrays = {} + for col in replacement_cols: + if col in transformed_cols_set: + updated_arrays[col] = transformed.arrow_table[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed_cols_set: + updated_arrays[unique_col] = transformed.arrow_table[col] + + updated_types = base.types.copy() + + for col in removed_cols: + if col in updated_types: + del updated_types[col] + + for col in replacement_cols: + if col in transformed.types: + updated_types[col] = transformed.types[col] + for col, unique_col in col_name_mapping.items(): + if col in transformed.types: + updated_types[unique_col] = transformed.types[col] + + modified_dataset = modify_table(base, updated_arrays, types=updated_types) + modified_dataset = modified_dataset.select_columns(new_columns_order) + + return modified_dataset + + +class ConverterScopeMixin: + """Scope handling shared by the converter units. + + Deliberately **not** named ``Base*`` and deliberately not inheriting from + ``BaseUnit``. ``ComponentRegistry._get_base_type`` walks the MRO for + ancestors whose name contains "Base" and that declare a ``TYPE``, and + rejects a component with more than one candidate — an intermediate class + called ``BaseConverterUnit`` would be a second candidate and break the + registration of every unit inheriting from it. + """ + + def _scope(self) -> Dict[str, List]: + """The configured scope, tolerating an explicit ``None``. + + The API schema always writes the ``scope`` key but allows it to be null, + so a converter saved without a scope arrives here as ``None`` rather + than as a missing key. + """ + return self.config.get("scope") or EMPTY_SCOPE + + def _target_index(self) -> Optional[int]: + """The 1-based target column index, or None when there is no target.""" + target = self.config.get("target") + if target is None: + return None + return target.get("idx") + + def _check_target_bounds(self, dataset: "DashAIDataset") -> None: + """Reject a target index that does not point at a column. + + Checked against the dataset in hand rather than once up front: when + several converters are chained the column count changes underneath, and + an unchecked index would surface as a bare IndexError. + """ + target_column_index = self._target_index() + if target_column_index is None: + return + + if int(target_column_index) < 1 or int(target_column_index) > len( + dataset.features + ): + raise JobError( + f"Target column index {target_column_index} is out of bounds" + ) + + def _resolve_scope( + self, dataset: "DashAIDataset" + ) -> Tuple[List[str], List[int], Optional[str]]: + """Turn the 1-based scope into names and indexes for **this** dataset. + + Always resolved against the dataset in hand, never against a list + carried through the context: a converter that renames, drops or adds + columns changes what index 3 means for whatever runs next. + """ + scope = self._scope() + column_names = dataset.column_names + + columns_scope = [column["idx"] - 1 for column in scope["columns"]] + scope_column_indexes = sorted(set(columns_scope)) + + if not scope_column_indexes: + scope_column_indexes = list(range(len(dataset.features))) + + scope_column_names = [column_names[index] for index in scope_column_indexes] + + rows_scope = [row - 1 for row in scope["rows"]] + scope_rows_indexes = sorted(set(rows_scope)) + + target_column_name = None + target_column_index = self._target_index() + if target_column_index is not None: + target_column_name = column_names[int(target_column_index) - 1] + + return scope_column_names, scope_rows_indexes, target_column_name + + def _slice_for_fit( + self, + dataset: "DashAIDataset", + scope_column_names: List[str], + scope_rows_indexes: List[int], + target_column_name: Optional[str], + ) -> Tuple["DashAIDataset", Optional["DashAIDataset"]]: + """The X and y a converter is fitted on: scoped columns, scoped rows.""" + y_dataset = None + if target_column_name is not None: + y_dataset = dataset.select_columns([target_column_name]) + if scope_rows_indexes: + y_dataset = y_dataset.select(scope_rows_indexes) + + x_dataset = dataset.select_columns(scope_column_names) + if scope_rows_indexes: + x_dataset = x_dataset.select(scope_rows_indexes) + + return x_dataset, y_dataset + + def _slice_for_transform( + self, + dataset: "DashAIDataset", + scope_column_names: List[str], + target_column_name: Optional[str], + ) -> Tuple["DashAIDataset", Optional["DashAIDataset"]]: + """The X and y a converter transforms: scoped columns, **every** row. + + A row scope narrows what the converter learns from, never what it is + applied to. + """ + y_dataset = None + if target_column_name is not None: + y_dataset = dataset.select_columns([target_column_name]) + + return dataset.select_columns(scope_column_names), y_dataset + + def _merge_transformed( + self, + converter_instance, + dataset: "DashAIDataset", + transformed_dataset: "DashAIDataset", + scope_column_names: List[str], + ) -> "DashAIDataset": + """Fold the transform output back into the dataset. + + A converter that changes the row count cannot have its columns merged + back position by position, so its output replaces the dataset outright. + """ + if type(converter_instance).CHANGES_ROW_COUNT: + return transformed_dataset + + return rebuild_dataset_with_transformed_columns( + dataset, + transformed_dataset, + scope_column_names, + ) + + def _fit(self, converter_instance, x_dataset, y_dataset, converter_name): + """Fit a converter, preserving the job's two error messages.""" + try: + return converter_instance.fit(x_dataset, y_dataset) + except ValueError as e: + log.error(f"Validation error in {converter_name}: {e}") + raise JobError(f"Validation error fitting {converter_name}: {e}") from e + except Exception as e: + log.exception(e) + raise JobError(f"Error fitting converter {converter_name}: {e}") from e + + def _transform(self, converter_instance, x_dataset, y_dataset, converter_name): + """Transform with a fitted converter, preserving the job's message.""" + try: + return converter_instance.transform(x_dataset, y_dataset) + except Exception as e: + log.exception(e) + raise JobError(f"Error transforming data with {converter_name}: {e}") from e diff --git a/DashAI/back/units/fit_converter_unit.py b/DashAI/back/units/fit_converter_unit.py new file mode 100644 index 000000000..645f10d52 --- /dev/null +++ b/DashAI/back/units/fit_converter_unit.py @@ -0,0 +1,107 @@ +"""Unit that fits a converter on a dataset without transforming anything.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import BaseSchema +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.converter_scope import ( + ConverterScopeMixin, + converter_field, + scope_field, + target_field, +) + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class FitConverterSchema(BaseSchema): + converter: converter_field() # type: ignore + scope: scope_field() # type: ignore + target: target_field() # type: ignore + + +class FitConverterUnit(BaseUnit, ConverterScopeMixin): + """Fit a converter on the dataset in the context and publish it fitted. + + Splits the "fit" half out of :class:`ApplyConverterUnit` so a converter can + be learned from one dataset and applied to another. That is the standard + train/test discipline: a scaler, an encoder or an imputer must learn its + statistics from the training data only, and then be applied unchanged to the + test data — refitting on test would leak information into the evaluation. + + The dataset is left untouched: this unit produces a fitted converter, not + data. Pair it with :class:`TransformDatasetUnit`, once per dataset the + fitted converter has to be applied to. + """ + + SCHEMA = FitConverterSchema + + REQUIRES = ("dataset",) + PROVIDES = ("fitted_converter",) + + def __init__(self, **config) -> None: + super().__init__(**config) + # Memoized on the instance, never in the context: two units of this + # class can live in the same context and they are different converters. + self._converter_class = None + + @property + def _converter_name(self) -> str: + return self.config["converter"]["component"] + + def _resolve_converter(self): + """Look the converter class up in the registry, once per instance.""" + if self._converter_class is not None: + return self._converter_class + + from kink import di + + component_registry = di["component_registry"] + converter_name = self._converter_name + + try: + self._converter_class = component_registry[converter_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Error importing converter {converter_name}: {e}") from e + + return self._converter_class + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + dataset: "DashAIDataset" = ctx.require("dataset") + converter_name = self._converter_name + + converter_constructor = self._resolve_converter() + converter_instance = converter_constructor( + **(self.config["converter"].get("params") or {}) + ) + + self._check_target_bounds(dataset) + ( + scope_column_names, + scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(dataset) + + log.info(f"Fitting converter: {converter_name}") + + x_dataset, y_dataset = self._slice_for_fit( + dataset, + scope_column_names, + scope_rows_indexes, + target_column_name, + ) + + fitted = self._fit(converter_instance, x_dataset, y_dataset, converter_name) + + ctx.put("fitted_converter", fitted) diff --git a/DashAI/back/units/transform_dataset_unit.py b/DashAI/back/units/transform_dataset_unit.py new file mode 100644 index 000000000..cba697af7 --- /dev/null +++ b/DashAI/back/units/transform_dataset_unit.py @@ -0,0 +1,82 @@ +"""Unit that transforms a dataset with an already fitted converter.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.converter_scope import ( + ConverterScopeMixin, + scope_field, + target_field, +) + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class TransformDatasetSchema(BaseSchema): + scope: scope_field() # type: ignore + target: target_field() # type: ignore + + +class TransformDatasetUnit(BaseUnit, ConverterScopeMixin): + """Apply an already fitted converter to the dataset in the context. + + Takes no converter configuration: the converter arrives fitted through the + context, from :class:`FitConverterUnit` or :class:`ApplyConverterUnit`. That + is what makes "fit on train, apply to test" possible without refitting — + run this once per dataset, and the converter keeps the state it learned. + + It does carry its own ``scope``, on purpose: the scope is a list of 1-based + indexes, and they are resolved against **this** unit's dataset. Reusing the + column names resolved during fit would break the moment the two datasets + order their columns differently, and would carry stale column identity + across the context boundary. + """ + + SCHEMA = TransformDatasetSchema + + REQUIRES = ("dataset", "fitted_converter") + PROVIDES = ("dataset",) + + def validate(self, ctx: ExecutionContext) -> None: + """Reject an out-of-bounds target before any work is done.""" + self._check_target_bounds(ctx.require("dataset")) + + def execute(self, ctx: ExecutionContext) -> None: + dataset: "DashAIDataset" = ctx.require("dataset") + converter_instance = ctx.require("fitted_converter") + converter_name = type(converter_instance).__name__ + + self._check_target_bounds(dataset) + ( + scope_column_names, + _scope_rows_indexes, + target_column_name, + ) = self._resolve_scope(dataset) + + log.info(f"Transforming with fitted converter: {converter_name}") + + x_dataset, y_dataset = self._slice_for_transform( + dataset, + scope_column_names, + target_column_name, + ) + + transformed_dataset = self._transform( + converter_instance, x_dataset, y_dataset, converter_name + ) + + ctx.put( + "dataset", + self._merge_transformed( + converter_instance, + dataset, + transformed_dataset, + scope_column_names, + ), + ) diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index 943662008..5286ab22a 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -11,6 +11,8 @@ "EvaluateModelUnit", "SaveModelUnit", "ApplyConverterUnit", + "FitConverterUnit", + "TransformDatasetUnit", "SaveDatasetUnit", } @@ -55,6 +57,16 @@ def test_unit_schemas_describe_their_configuration(units): "scope", "target", } + assert set(units["FitConverterUnit"]["schema"]["properties"]) == { + "converter", + "scope", + "target", + } + # No converter to pick: it arrives already fitted through the context. + assert set(units["TransformDatasetUnit"]["schema"]["properties"]) == { + "scope", + "target", + } # SaveDatasetUnit is configuration-free: it saves where the load said. assert units["SaveDatasetUnit"]["schema"]["properties"] == {} diff --git a/tests/back/units/test_converter_fit_transform_split.py b/tests/back/units/test_converter_fit_transform_split.py new file mode 100644 index 000000000..00cba74a1 --- /dev/null +++ b/tests/back/units/test_converter_fit_transform_split.py @@ -0,0 +1,278 @@ +"""Fitting a converter on one dataset and applying it to another. + +The point of splitting ``ApplyConverterUnit`` into ``FitConverterUnit`` + +``TransformDatasetUnit``: a scaler, encoder or imputer must learn its statistics +from the training data only and then be applied unchanged to the test data. +Refitting on test would leak the test distribution into the evaluation, which is +exactly what the fused unit forced. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.types.value_types import Float +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit + +FULL_SCOPE = {"columns": [], "rows": []} + + +def _dataset(**columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def _values(dataset, column="a"): + return list(dataset.to_pandas()[column]) + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = {"MinMaxScaler": {"class": MinMaxScaler}} + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def test_a_converter_fitted_on_train_is_applied_to_test_without_refitting(registry): + """The headline case. + + MinMaxScaler fitted on [0, 5, 10] learns min=0, max=10. Applied to a test + value of 20 it must yield 2.0 — outside [0, 1] precisely because the range + came from train. A refit on the test data would have produced 0.0 instead, + so the number is what proves the fitted state survived. + """ + ctx = ExecutionContext() + + # Fit on train. + ctx.put("dataset", _dataset(a=[0.0, 5.0, 10.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + # Transform train with it. + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + assert _values(ctx.require("dataset")) == [0.0, 0.5, 1.0] + + # Swap in the test dataset and transform with the *same* fitted converter. + ctx.put("dataset", _dataset(a=[20.0])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [2.0] + + +def test_fitting_leaves_the_dataset_untouched(registry): + """FitConverterUnit produces a converter, not data.""" + ctx = ExecutionContext() + original = _dataset(a=[0.0, 5.0, 10.0]) + ctx.put("dataset", original) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + assert ctx.require("dataset") is original + assert _values(ctx.require("dataset")) == [0.0, 5.0, 10.0] + + +def test_the_fitted_converter_is_published_live_not_copied(registry): + """It has to be the same object, or the learned state would be lost.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + fitted = ctx.require("fitted_converter") + assert isinstance(fitted, MinMaxScaler) + assert ctx.require("fitted_converter") is fitted + # The learned statistics are what makes it worth reusing. + assert list(fitted.data_min_) == [0.0] + assert list(fitted.data_max_) == [10.0] + + +def test_transforming_without_a_fitted_converter_is_a_contract_error(): + """A missing converter is a wiring mistake, not "nothing to apply".""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0])) + + with pytest.raises(UnitContractError, match="'fitted_converter' is not available"): + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + +def test_transforming_without_a_dataset_is_a_contract_error(registry): + ctx = ExecutionContext() + ctx.put("fitted_converter", MinMaxScaler()) + + with pytest.raises(UnitContractError, match="'dataset' is not available"): + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + +def test_the_fused_unit_also_publishes_its_fitted_converter(registry): + """ApplyConverterUnit stays usable as the source of a reusable converter. + + So the single-dataset path and the train/test path are the same mechanism: + whoever fitted the converter publishes it, and any number of + TransformDatasetUnits can then apply it elsewhere. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 5.0, 10.0])) + + ApplyConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + assert _values(ctx.require("dataset")) == [0.0, 0.5, 1.0] + + ctx.put("dataset", _dataset(a=[20.0])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [2.0] + + +def test_the_split_pair_matches_the_fused_unit_on_a_single_dataset(registry): + """The two paths must not drift: same input, same output.""" + fused_ctx = ExecutionContext() + fused_ctx.put("dataset", _dataset(a=[1.0, 2.0, 3.0, 4.0])) + ApplyConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(fused_ctx) + + split_ctx = ExecutionContext() + split_ctx.put("dataset", _dataset(a=[1.0, 2.0, 3.0, 4.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(split_ctx) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(split_ctx) + + assert _values(fused_ctx.require("dataset")) == _values( + split_ctx.require("dataset") + ) + + +def test_a_row_scope_narrows_the_fit_but_the_transform_still_sees_every_row(registry): + """Row scope belongs to fitting; transform always covers the dataset. + + Fitting on rows 1-2 of [0, 10, 100] learns min=0, max=10, so the third row + scales to 10.0 rather than 1.0. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0, 100.0])) + + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope={"columns": [], "rows": [1, 2]}, + target=None, + )(ctx) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + + assert _values(ctx.require("dataset")) == [0.0, 1.0, 10.0] + + +def test_fit_names_the_culprit_when_the_converter_is_not_registered(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0])) + + with pytest.raises(JobError, match="Error importing converter NotRegistered"): + FitConverterUnit( + converter={"component": "NotRegistered", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + +def test_fit_rejects_an_out_of_bounds_target_before_running(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1.0], b=[2.0])) + unit = FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target={"idx": 99}, + ) + + with pytest.raises(JobError, match="Target column index 99 is out of bounds"): + unit.validate(ctx) + + assert not ctx.has("fitted_converter") + + +def test_two_fit_units_in_one_context_resolve_their_converters_independently(registry): + """Registry lookups are memoized on the instance, never in the context.""" + + class _Other(MinMaxScaler): + pass + + registry["Other"] = {"class": _Other} + + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + + first = FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + ) + second = FitConverterUnit( + converter={"component": "Other", "params": {}}, + scope=FULL_SCOPE, + target=None, + ) + first(ctx) + second(ctx) + + assert first._converter_class is MinMaxScaler + assert second._converter_class is _Other + + +def test_the_units_register_under_the_unit_type(): + """Guards the mixin's name. + + ``ConverterScopeMixin`` must not be called ``Base*``: the registry rejects a + component whose MRO has more than one "Base" ancestor declaring a TYPE, so + an intermediate ``BaseConverterUnit`` would break registration for all three + converter units at once. + """ + from DashAI.back.dependencies.registry import ComponentRegistry + + units = [ApplyConverterUnit, FitConverterUnit, TransformDatasetUnit] + component_registry = ComponentRegistry(initial_components=units) + + for unit in units: + assert component_registry[unit.__name__]["type"] == "Unit" + + +def test_one_fitted_converter_feeds_several_transforms(registry): + """Nothing about a transform consumes or invalidates the fitted converter.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[0.0, 10.0])) + FitConverterUnit( + converter={"component": "MinMaxScaler", "params": {}}, + scope=FULL_SCOPE, + target=None, + )(ctx) + + for value, expected in ((5.0, 0.5), (20.0, 2.0), (-10.0, -1.0)): + ctx.put("dataset", _dataset(a=[value])) + TransformDatasetUnit(scope=FULL_SCOPE, target=None)(ctx) + assert _values(ctx.require("dataset")) == [expected] From d8f5c4c9b053bac3f7b416c0a2cbe7814e8b42b1 Mon Sep 17 00:00:00 2001 From: Felipe Date: Wed, 5 Aug 2026 16:38:59 -0400 Subject: [PATCH 06/28] test: Add the over-declared REQUIRES check to the contract audit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Brings the fourth contract check from feat/atom-expl-predict-explor, where it was written, so the file is byte-identical on both branches and the two PRs cannot diverge on it. The check matters on its own: __call__ demands every key in REQUIRES unconditionally, so a key that is declared but never read is not harmless documentation — it rejects any upstream that does not happen to publish it. It also matters that this file specifically stays in sync. A mismerge here is the one that fails silently: the audit keeps passing, it just audits less. Co-Authored-By: Claude Opus 5 (1M context) --- tests/back/units/test_unit_contracts.py | 26 +++++++++++++++++++++++++ 1 file changed, 26 insertions(+) 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 946aa3a94cf65d9609cce409f2897c9faf799ba2 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 3 Aug 2026 19:11:02 -0400 Subject: [PATCH 07/28] 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 | 146 ++-- 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 +++++++++++ 30 files changed, 5419 insertions(+), 577 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 4f3870264..ee57db1c5 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 matrices, so the background must @@ -255,7 +258,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict_prepared, + 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 539b23482..add16113b 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_prepared, + 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 bc84d0871..4f2f2904e 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 matrices, so the background must @@ -196,7 +199,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict_prepared, + 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 3126ed3b5..328b34575 100644 --- a/DashAI/back/explainability/model_input.py +++ b/DashAI/back/explainability/model_input.py @@ -17,12 +17,50 @@ the already prepared matrix a second time. """ -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Callable, List, Optional 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 db86e671f..06b815725 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -348,14 +348,31 @@ # 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_converter_unit import FitConverterUnit 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 from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit logging.basicConfig(level=logging.DEBUG) @@ -531,6 +548,19 @@ def get_initial_components(): FitConverterUnit, TransformDatasetUnit, 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 beb261e95..d4643bd0b 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,18 +88,15 @@ 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 from DashAI.back.exploration.artifact_store import store_artifacts - 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: @@ -124,103 +124,61 @@ 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) - explorer_info.set_status_as_error() - db.commit() - raise JobError( - ( - f"Explorer {explorer_info.exploration_type} " - "not found in the registry." - ) - ) from e - - # Instance the explorer (the explorer handles its validation) - try: - explorer_instance = explorer_component_class(**explorer_info.parameters) - assert isinstance(explorer_instance, BaseExplorer) - except Exception 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"Error instancing the explorer {explorer_info.exploration_type}." - ) from e + raise - # prepare the dataset + # Load the dataset from the notebook: its own working copy, which + # is what the converters rewrite. try: - prepared_dataset = explorer_instance.prepare_dataset( - loaded_dataset, explorer_info.columns - ) + 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( - ( - "Error preparing the dataset for the exploration " - f"{explorer_info.exploration_type}." - ) - ) from e - - # Launch the exploration + 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: - result = explorer_instance.launch_exploration( - prepared_dataset, explorer_info - ) + 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( - 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() + # Update the explorer info. The status is not set to finished + # here: the artifacts below are part of the work, so the row + # only counts as done once they exist too. + explorer_info.exploration_path = ctx.require("exploration_path") db.commit() except Exception as e: log.exception(e) @@ -237,9 +195,17 @@ def run( # the explorer class is asked for its results: from here on the # stored artifacts are served as is, so the exploration keeps # rendering even if the explorer is removed from the registry. + # + # Both inputs come from the context rather than from local + # variables: the explorer instance is what ran the exploration + # (published by RunExplorationUnit, so the artifacts are built from + # the same object that produced the result, not a rebuilt one), and + # the path is where SaveExplorationUnit actually wrote it. try: explorer_info.artifacts_path = store_artifacts( - explorer_instance, save_path, explorer_info.id + ctx.require("explorer"), + ctx.require("exploration_path"), + explorer_info.id, ) explorer_info.set_status_as_finished() db.commit() @@ -253,3 +219,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 5286ab22a..6443545d1 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -14,6 +14,19 @@ "FitConverterUnit", "TransformDatasetUnit", "SaveDatasetUnit", + "RunExplorationUnit", + "SaveExplorationUnit", + "LoadTrainedModelUnit", + "LoadTrainingDatasetUnit", + "BuildManualInputUnit", + "PredictUnit", + "SavePredictionUnit", + "LoadRunModelUnit", + "BuildGlobalExplainerUnit", + "BuildLocalExplainerUnit", + "PrepareExplanationDataUnit", + "GenerateGlobalExplanationUnit", + "GenerateLocalExplanationUnit", } @@ -69,6 +82,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): @@ -81,10 +120,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) From 215a813c3392b1546c07c3e9d60c0e10d3e5a055 Mon Sep 17 00:00:00 2001 From: Felipe Date: Wed, 5 Aug 2026 17:38:39 -0400 Subject: [PATCH 08/28] 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: ") From c043515a09e5c14fe63d425b0f4bfdfb002ae547 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 10 Aug 2026 01:20:42 -0400 Subject: [PATCH 09/28] fix test python 3.10 --- tests/back/explainers/test_shap_predictor_handover.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/back/explainers/test_shap_predictor_handover.py b/tests/back/explainers/test_shap_predictor_handover.py index b5bda46ed..fb43ce6b2 100644 --- a/tests/back/explainers/test_shap_predictor_handover.py +++ b/tests/back/explainers/test_shap_predictor_handover.py @@ -12,6 +12,9 @@ AttributeError: property 'feature_names_in_' of 'LGBMClassifier' object has no setter +(on Python 3.10 the same failure reads ``can't set attribute +'feature_names_in_'``). + ``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 @@ -53,7 +56,8 @@ def test_these_models_really_do_expose_feature_names_read_only(model_class, fram model.fit(x, y) assert hasattr(model, "feature_names_in_") - with pytest.raises(AttributeError, match="no setter"): + # CPython worded this differently before 3.11 ("can't set attribute"). + with pytest.raises(AttributeError, match="no setter|can't set attribute"): model.feature_names_in_ = None From cb91ff87ac511d7e4ace564810dbab999e091da6 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 10 Aug 2026 02:01:10 -0400 Subject: [PATCH 10/28] remove lightgbm xgboost --- DashAI/back/explainability/model_input.py | 10 +- .../test_shap_predictor_handover.py | 119 +++++------------- 2 files changed, 37 insertions(+), 92 deletions(-) diff --git a/DashAI/back/explainability/model_input.py b/DashAI/back/explainability/model_input.py index 328b34575..2580345b2 100644 --- a/DashAI/back/explainability/model_input.py +++ b/DashAI/back/explainability/model_input.py @@ -32,11 +32,11 @@ def as_shap_predictor(model: Any) -> Callable: 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. + That assumption does not hold for estimators that expose + ``feature_names_in_`` as a read-only ``property`` — a common shape among + third party wrappers — where 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 diff --git a/tests/back/explainers/test_shap_predictor_handover.py b/tests/back/explainers/test_shap_predictor_handover.py index fb43ce6b2..230f2eef4 100644 --- a/tests/back/explainers/test_shap_predictor_handover.py +++ b/tests/back/explainers/test_shap_predictor_handover.py @@ -3,99 +3,17 @@ ``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. +attribute is writable, which does not hold for every estimator: those that +expose ``feature_names_in_`` as a read-only ``property`` raise instead, and the +explanation dies before it starts. -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 - -(on Python 3.10 the same failure reads ``can't set attribute -'feature_names_in_'``). - -``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``. +``as_shap_predictor`` hands over a plain closure, which has no ``__self__``, so +SHAP skips that step entirely. These tests pin the mechanism. """ -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_") - # CPython worded this differently before 3.11 ("can't set attribute"). - with pytest.raises(AttributeError, match="no setter|can't set attribute"): - 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(): @@ -122,3 +40,30 @@ def predict(self, x): model = Model() assert getattr(model.predict, "__self__", None) is model assert getattr(as_shap_predictor(model), "__self__", None) is None + + +def test_the_wrapped_predictor_survives_shaps_model_conversion(): + """``convert_to_model`` must accept the closure without touching a model.""" + from shap.utils._legacy import convert_to_model + + class ReadOnlyFeatureNames: + """Stands in for an estimator whose ``feature_names_in_`` has no setter.""" + + @property + def feature_names_in_(self): + return ["a", "b"] + + def predict(self, x): + return [0] + + model = ReadOnlyFeatureNames() + + # The failure this exists to prevent: SHAP reaches the model through + # ``__self__`` and tries to blank the attribute. + with pytest.raises(AttributeError, match="feature_names_in_"): + convert_to_model(model.predict) + + converted = convert_to_model(as_shap_predictor(model)) + + assert converted.f is not None + assert list(model.feature_names_in_) == ["a", "b"] From df47bc3694d09444fd17651bfc10c7f793b908f8 Mon Sep 17 00:00:00 2001 From: Felipe Date: Thu, 13 Aug 2026 21:26:56 -0400 Subject: [PATCH 11/28] fix: route the SHAP predictor through predict_prepared as_shap_predictor handed SHAP a closure over model.predict. The callers (KernelShap, RegressionKernelShap, ContrastiveShap) first move the background into the model's feature space with prepare_model_input, so going through predict ran the model's input preparation a second time over an already prepared matrix. SHAP then perturbs that matrix into plain arrays, which the preparation cannot consume at all, and the failure surfaced far from its cause as AttributeError: 'numpy.ndarray' object has no attribute 'types' The module docstring and a comment above each caller already said the model had to be queried through predict_prepared; only the call was left behind. test_the_wrapped_predictor_never_routes_through_predict pins it: its stub raises if predict is reached, so a future regression fails at the wrapper instead of five frames away inside SHAP. --- DashAI/back/explainability/model_input.py | 15 ++++++--- .../test_shap_predictor_handover.py | 31 +++++++++++++++---- 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/DashAI/back/explainability/model_input.py b/DashAI/back/explainability/model_input.py index 2580345b2..c459f1018 100644 --- a/DashAI/back/explainability/model_input.py +++ b/DashAI/back/explainability/model_input.py @@ -24,7 +24,7 @@ def as_shap_predictor(model: Any) -> Callable: - """Wrap ``model.predict`` so SHAP receives a plain function, not a method. + """Wrap the model's prepared-matrix prediction so SHAP gets a plain function. SHAP suppresses scikit-learn's "X does not have valid feature names" warning by blanking ``feature_names_in_`` on whatever object the callable @@ -43,6 +43,13 @@ def as_shap_predictor(model: Any) -> Callable: interface for ``model``. The only thing lost is the suppression of a cosmetic scikit-learn warning. + It routes to ``predict_prepared``, not to ``predict``. Callers hand SHAP a + background already moved into the model's feature space with + ``prepare_model_input``, and SHAP then queries the model with perturbed + copies of *that* matrix. Going through ``predict`` would run the model's + input preparation a second time over an already prepared matrix — and SHAP + passes plain arrays, which the preparation cannot consume at all. + Parameters ---------- model : Any @@ -51,12 +58,12 @@ def as_shap_predictor(model: Any) -> Callable: Returns ------- Callable - A one-argument function calling ``model.predict`` positionally, the - same way SHAP calls it today. + A one-argument function calling ``model.predict_prepared`` + positionally, the same way SHAP calls it. """ def predict(x): - return model.predict(x) + return model.predict_prepared(x) return predict diff --git a/tests/back/explainers/test_shap_predictor_handover.py b/tests/back/explainers/test_shap_predictor_handover.py index 230f2eef4..dcce6df25 100644 --- a/tests/back/explainers/test_shap_predictor_handover.py +++ b/tests/back/explainers/test_shap_predictor_handover.py @@ -16,12 +16,12 @@ from DashAI.back.explainability.model_input import as_shap_predictor -def test_the_wrapped_predictor_forwards_to_predict_positionally(): +def test_the_wrapped_predictor_forwards_to_predict_prepared_positionally(): """SHAP calls the model with one positional argument; that must not change.""" seen = {} class Model: - def predict(self, x): + def predict_prepared(self, x): seen["arg"] = x return [0] @@ -30,15 +30,34 @@ def predict(self, x): assert seen["arg"] == "the frame" +def test_the_wrapped_predictor_never_routes_through_predict(): + """The callers hand SHAP a background already in the model's feature space. + + Routing to ``predict`` would prepare an already prepared matrix a second + time, and SHAP perturbs it into plain arrays that the preparation cannot + consume at all — which surfaces far away, as + ``'numpy.ndarray' object has no attribute 'types'``. + """ + + class Model: + def predict(self, x): + raise AssertionError("predict would prepare an already prepared matrix") + + def predict_prepared(self, x): + return [1] + + assert as_shap_predictor(Model())("the frame") == [1] + + 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): + def predict_prepared(self, x): return [0] model = Model() - assert getattr(model.predict, "__self__", None) is model + assert getattr(model.predict_prepared, "__self__", None) is model assert getattr(as_shap_predictor(model), "__self__", None) is None @@ -53,7 +72,7 @@ class ReadOnlyFeatureNames: def feature_names_in_(self): return ["a", "b"] - def predict(self, x): + def predict_prepared(self, x): return [0] model = ReadOnlyFeatureNames() @@ -61,7 +80,7 @@ def predict(self, x): # The failure this exists to prevent: SHAP reaches the model through # ``__self__`` and tries to blank the attribute. with pytest.raises(AttributeError, match="feature_names_in_"): - convert_to_model(model.predict) + convert_to_model(model.predict_prepared) converted = convert_to_model(as_shap_predictor(model)) From ce22e7dafa7b80313182be81d0def0bd9afa04f5 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 10 Aug 2026 01:16:14 -0400 Subject: [PATCH 12/28] Add contract tests for dataset ingestion units and update unit schemas - Introduced new tests for `LoadUploadedDatasetUnit`, `LoadDatafileDatasetUnit`, `InferDatasetTypesUnit`, `ApplyDatasetSchemaUnit`, `ComputeDatasetMetadataUnit`, and `SaveDatasetToPathUnit` in `test_dataset_ingest_units.py`. - Updated expected unit schemas in `test_units_api.py` to include new dataset ingestion units. - Enhanced validation checks and error handling in the dataset processing workflow. --- DashAI/back/dependencies/database/models.py | 7 +- DashAI/back/initial_components.py | 12 + DashAI/back/job/dataset_job.py | 279 ++--- .../back/units/apply_dataset_schema_unit.py | 129 ++ .../units/compute_dataset_metadata_unit.py | 136 ++ DashAI/back/units/infer_dataset_types_unit.py | 76 ++ .../back/units/load_datafile_dataset_unit.py | 178 +++ .../back/units/load_uploaded_dataset_unit.py | 154 +++ .../back/units/save_dataset_to_path_unit.py | 78 ++ DashAI/back/units/save_dataset_unit.py | 8 +- tests/back/api/test_dataset_job.py | 1094 +++++++++++++++++ tests/back/api/test_units_api.py | 37 + tests/back/units/test_dataset_ingest_units.py | 491 ++++++++ 13 files changed, 2517 insertions(+), 162 deletions(-) create mode 100644 DashAI/back/units/apply_dataset_schema_unit.py create mode 100644 DashAI/back/units/compute_dataset_metadata_unit.py create mode 100644 DashAI/back/units/infer_dataset_types_unit.py create mode 100644 DashAI/back/units/load_datafile_dataset_unit.py create mode 100644 DashAI/back/units/load_uploaded_dataset_unit.py create mode 100644 DashAI/back/units/save_dataset_to_path_unit.py create mode 100644 tests/back/api/test_dataset_job.py create mode 100644 tests/back/units/test_dataset_ingest_units.py diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py index 761c76f95..07741791c 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -106,10 +106,14 @@ def set_status_as_delivered(self) -> None: def set_status_as_started(self) -> None: """ Update the status of the dataset to started and set created to now. + + Unlike ``Run`` or ``Explorer``, this table has no ``start_time`` / + ``end_time`` columns, so there is nothing to stamp beyond the timestamps + it does declare. Assigning them anyway only set an attribute on the + instance that was dropped at commit and read back as missing. """ self.status = DatasetStatus.STARTED self.created = datetime.now() - self.start_time = datetime.now() def set_status_as_finished(self) -> None: """ @@ -117,7 +121,6 @@ def set_status_as_finished(self) -> None: """ self.status = DatasetStatus.FINISHED self.last_modified = datetime.now() - self.end_time = datetime.now() def set_status_as_error(self) -> None: """ diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 06b815725..bd507979b 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -348,10 +348,12 @@ # Units from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit 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.compute_dataset_metadata_unit import ComputeDatasetMetadataUnit from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_converter_unit import FitConverterUnit from DashAI.back.units.fit_model_unit import FitModelUnit @@ -361,14 +363,18 @@ from DashAI.back.units.generate_local_explanation_unit import ( GenerateLocalExplanationUnit, ) +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit 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.load_uploaded_dataset_unit import LoadUploadedDatasetUnit 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_to_path_unit import SaveDatasetToPathUnit 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 @@ -561,6 +567,12 @@ def get_initial_components(): PrepareExplanationDataUnit, GenerateGlobalExplanationUnit, GenerateLocalExplanationUnit, + LoadUploadedDatasetUnit, + LoadDatafileDatasetUnit, + InferDatasetTypesUnit, + ApplyDatasetSchemaUnit, + ComputeDatasetMetadataUnit, + SaveDatasetToPathUnit, # Explainers ContrastiveShap, DiceCounterfactual, diff --git a/DashAI/back/job/dataset_job.py b/DashAI/back/job/dataset_job.py index d7caef39a..cdabf63fe 100644 --- a/DashAI/back/job/dataset_job.py +++ b/DashAI/back/job/dataset_job.py @@ -8,6 +8,14 @@ from DashAI.back.api.utils import parse_params from DashAI.back.dependencies.database.models import Converter, Dataset, Notebook from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit +from DashAI.back.units.compute_dataset_metadata_unit import ComputeDatasetMetadataUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -98,14 +106,6 @@ def run( import uuid from pathlib import Path - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - transform_dataset_with_schema, - ) - from DashAI.back.types.inf.type_inference import infer_types - - component_registry = di["component_registry"] session_factory = di["session_factory"] config = di["config"] @@ -122,6 +122,8 @@ def run( temp_dir = tempfile.mkdtemp(prefix="dashai-dataset-") url = self.kwargs.get("url", "") + ctx = ExecutionContext() + try: with session_factory() as db: dataset = db.get(Dataset, dataset_id) @@ -134,6 +136,13 @@ def run( self.report_progress(0.1, "Loading data") + # Whether the destination folder is this job's to delete. Re-importing + # into an existing dataset writes over its current folder, and the + # failure paths below clean up by removing it — which would destroy + # data the surviving row still points at. Only a folder this run + # created may be removed. + folder_is_ours = False + if n_sample and dataset.file_path != "": folder_path = Path(dataset.file_path) else: @@ -148,6 +157,7 @@ def run( raise JobError( f"A dataset with the name {random_name} already exists." ) from e + folder_is_ours = True from_notebook_no_converters = False try: @@ -159,6 +169,9 @@ def run( .filter(Notebook.id == notebook_id) .first() ) + # Checked here rather than left to the load unit, whose + # own wording for a missing notebook differs from this + # one. The message reaches the UI, so it is preserved. if not notebook_dataset: msg = ( "Notebook with ID " @@ -178,208 +191,156 @@ def run( is not None ) from_notebook_no_converters = not has_converters - new_dataset = load_dataset( - os.path.join(notebook_dataset.file_path, "dataset") - ) + + # ``LoadDatasetUnit`` also publishes ``dataset_path`` (the + # notebook's own copy) and ``dataset_id`` (the *source* + # dataset). Neither describes what is being created here: the + # save goes to a new folder, and this job's ``dataset_id`` is + # the destination row, read from kwargs. Nothing below reads + # them from the context, and nothing should start to. + LoadDatasetUnit(notebook_id=notebook_id)(ctx) + + # No schema is applied to a notebook copy: it is already a + # stored dataset, with its types settled when it was created. else: source_name = self.kwargs.get("source_name") if source_name: # --- Hub import path --- - from DashAI.back.core.enums.status import DatafileStatus - from DashAI.back.dependencies.database.models import ( - Datafile, - ) - + # The id is required, and it is the request that is + # malformed without it, so the check stays here rather + # than becoming a unit that cannot be configured. datafile_id = params.get("datafile_id") - selected_file = params.get("selected_file") - if datafile_id is None: raise JobError("datafile_id is required for hub imports.") - with session_factory() as db: - hub_row = db.get(Datafile, datafile_id) - if hub_row is None or hub_row.status != DatafileStatus.READY: - raise JobError(f"Datafile {datafile_id} is not ready.") - hub_work_dir = hub_row.local_path - if selected_file: - file_path_hub = str(Path(hub_work_dir) / selected_file) - else: - hub_base = Path(hub_work_dir) - files = sorted( - str(p) - for p in hub_base.rglob("*") - if p.is_file() - and not any( - part.startswith(".") - for part in p.relative_to(hub_base).parts - ) - ) - if not files: - raise JobError("Hub download directory is empty.") - file_path_hub = files[0] - - selected_dataloader = params.get("dataloader", "") - _reg = component_registry._registry - dl_registry = _reg.get("DataLoader", {}) - if selected_dataloader not in dl_registry: - raise JobError( - f"DataLoader '{selected_dataloader}'" - " not found in registry." - ) - dataloader = dl_registry[selected_dataloader]["class"]() - log.debug( - "Loading hub dataset from %s using %s", - file_path_hub, - selected_dataloader, - ) - hub_loader_params = params.get("dataloader_params", {}) - new_dataset = dataloader.load_data( - filepath_or_buffer=file_path_hub, - temp_path=hub_work_dir, - params=hub_loader_params, - n_sample=None, - ) + LoadDatafileDatasetUnit( + dataloader={ + "component": params.get("dataloader", ""), + "params": params.get("dataloader_params", {}), + }, + datafile_id=datafile_id, + selected_file=params.get("selected_file"), + )(ctx) else: - # --- File / URL upload path (unchanged) --- + # --- File / URL upload path --- + # Validating the request's params is unpacking of how the + # upload arrived, not part of reading the data: the unit + # gets the reader already picked and its params already + # checked. ``model_dump()`` keeps the payload the reader + # receives exactly as it was. parsed_params = parse_params(DatasetParams, json.dumps(params)) - dataloader = component_registry[parsed_params.dataloader][ - "class" - ]() log.debug("Storing dataset in %s", folder_path) - new_dataset = dataloader.load_data( - filepath_or_buffer=( - str(file_path) if file_path is not None else url - ), + LoadUploadedDatasetUnit( + dataloader={ + "component": parsed_params.dataloader, + "params": parsed_params.model_dump(), + }, + source=str(file_path) if file_path is not None else url, temp_path=str(temp_dir), - params=parsed_params.model_dump(), n_sample=n_sample, - ) + )(ctx) + # The types either come with the request or are worked out + # from the data. Both paths end in the same context key, so + # the unit that applies them has a single input either way. if params.get("inferred_types"): - schema = params["inferred_types"] - elif new_dataset.types: - schema = { - col: typ.to_string() - for col, typ in new_dataset.types.items() - } + ctx.put_ref("inferred_types", params["inferred_types"]) else: - schema = infer_types( - new_dataset.to_pandas(), method="DashAIPtype" - ) - if "column_renames" in params: - renames = params["column_renames"] - original_names = new_dataset.arrow_table.schema.names - new_names = [renames.get(col, col) for col in original_names] - - if len(new_names) != len(set(new_names)): - duplicate_names = set() - seen = set() - for name in new_names: - if name in seen: - duplicate_names.add(name) - else: - seen.add(name) - msg = ( - "Invalid column_renames: resulting column names " - "contain duplicates: " - f"{sorted(duplicate_names)}" - ) - raise JobError(msg) - - arrow_table = new_dataset.arrow_table.rename_columns(new_names) - new_dataset = new_dataset.__class__( - arrow_table, - splits=new_dataset.splits, - types=new_dataset.types, - ) - schema = {renames.get(col, col): schema[col] for col in schema} + InferDatasetTypesUnit(method="DashAIPtype")(ctx) - new_dataset = transform_dataset_with_schema(new_dataset, schema) + ApplyDatasetSchemaUnit( + column_renames=params.get("column_renames"), + )(ctx) self.report_progress(0.5, "Computing metadata") - compute_meta = params.get("compute_metadata", True) - extended_keys = ( - "general_info", - "numeric_stats", - "categorical_stats", - "text_stats", - "quality_info", - "correlations", - ) - - if from_notebook_no_converters: - # No converters applied - saved data matches the source - # dataset byte-for-byte. Reuse the source's splits.json - # (already loaded into ``new_dataset.splits`` by - # ``load_dataset``) instead of recomputing. - has_extended = any(k in new_dataset.splits for k in extended_keys) - if compute_meta: - # Use source's full metadata if present, else compute. - if not has_extended: - new_dataset.compute_metadata() - else: - # Keep only base metadata; drop any inherited extended. - if "total_rows" not in new_dataset.splits: - new_dataset.compute_base_metadata() - for stale_key in extended_keys: - new_dataset.splits.pop(stale_key, None) - elif compute_meta: - new_dataset.compute_metadata() - else: - new_dataset.compute_base_metadata() - # Defensive: strip any extended keys that may have been - # inherited from a source dataset (e.g. notebook flow - # with converters that ran before this rule existed). - for stale_key in extended_keys: - new_dataset.splits.pop(stale_key, None) + # ``from_notebook_no_converters`` means the saved data matches + # the source dataset byte-for-byte, so the metadata it arrived + # with still describes it. That is a policy decision this job + # makes from the notebook's converter history; the unit only + # needs the answer. + ComputeDatasetMetadataUnit( + compute_metadata=params.get("compute_metadata", True), + trust_inherited_metadata=from_notebook_no_converters, + )(ctx) gc.collect() self.report_progress(0.8, "Saving dataset") dataset_save_path = folder_path / "dataset" log.debug("Saving dataset in %s", str(dataset_save_path)) - save_dataset(new_dataset, dataset_save_path) + SaveDatasetToPathUnit(path=str(dataset_save_path))(ctx) except Exception as e: log.exception(e) - shutil.rmtree(folder_path, ignore_errors=True) + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) raise JobError(f"Error loading dataset: {str(e)}") from e - # Add dataset to database + # Add dataset to database. The counts are read back off the dataset + # rather than published by the metadata unit: they describe the + # dataset as it is right now, and a key holding them would go stale + # the moment anything else transformed it. + stored_metadata = ctx.require("dataset").splits with session_factory() as db: log.debug("Storing dataset metadata in database.") try: folder_path = os.path.realpath(folder_path) dataset = db.get(Dataset, dataset_id) + # Re-read, so it can be gone by now: the row is deletable + # through the API while the job runs. Without this guard the + # assignment below raises AttributeError on None, and the + # data just written to disk is orphaned. + if dataset is None: + raise JobError( + f"Dataset with ID {dataset_id} no longer exists." + ) dataset.file_path = folder_path - dataset.total_rows = new_dataset.splits.get("total_rows") - dataset.total_columns = len( - new_dataset.splits.get("column_names", []) - ) + dataset.total_rows = stored_metadata.get("total_rows") + dataset.total_columns = len(stored_metadata.get("column_names", [])) dataset.set_status_as_finished() db.commit() db.refresh(dataset) except exc.SQLAlchemyError as e: log.exception(e) - shutil.rmtree(folder_path, ignore_errors=True) + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) raise JobError("Internal database error") from e + except Exception: + # Anything else here leaves a dataset on disk that no row + # points at, so it has to be cleaned up too. Re-raised + # unchanged; the handler below writes the status. + if folder_is_ours: + shutil.rmtree(folder_path, ignore_errors=True) + raise log.debug("Dataset creation successfully finished.") - except JobError as e: + except Exception as e: + # Every failure, not just ``JobError``. Nothing else moves this row + # out of STARTED: Huey's error signal writes to its own ``task_copy`` + # table and never touches ``Dataset``, so an exception this handler + # does not see leaves the dataset stuck as in-progress forever, and + # the UI keeps showing a spinner for work that already died. log.error(f"Dataset creation failed: {e}") - with session_factory() as db: - dataset = db.get(Dataset, dataset_id) - if dataset: - dataset.set_status_as_error() - db.commit() - db.refresh(dataset) - raise e + try: + with session_factory() as db: + dataset = db.get(Dataset, dataset_id) + if dataset: + dataset.set_status_as_error() + db.commit() + db.refresh(dataset) + except Exception as bookkeeping_error: + # Never let the bookkeeping mask what actually went wrong. + log.exception(bookkeeping_error) + # Re-raised as-is: the message is the contract, and wrapping it here + # would change what every branch above reports. + raise finally: + ctx.clear_cache() gc.collect() if temp_dir and os.path.exists(temp_dir): try: diff --git a/DashAI/back/units/apply_dataset_schema_unit.py b/DashAI/back/units/apply_dataset_schema_unit.py new file mode 100644 index 000000000..42f60af28 --- /dev/null +++ b/DashAI/back/units/apply_dataset_schema_unit.py @@ -0,0 +1,129 @@ +"""Unit that renames columns and casts them to their declared types.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + none_type, + 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 + +log = logging.getLogger(__name__) + + +class ApplyDatasetSchemaSchema(BaseSchema): + column_renames: schema_field( + none_type(dict), + placeholder=None, + description=MultilingualString( + en="Mapping of current column name to new name, applied before the " + "types are cast. Columns not mentioned keep their name. The " + "resulting names have to stay unique.", + es="Mapa de nombre de columna actual a nombre nuevo, aplicado antes " + "de convertir los tipos. Las columnas no mencionadas conservan su " + "nombre. Los nombres resultantes tienen que seguir siendo únicos.", + pt="Mapa de nome de coluna atual para nome novo, aplicado antes de " + "converter os tipos. As colunas não mencionadas mantêm o seu nome. " + "Os nomes resultantes têm de continuar únicos.", + de="Zuordnung von aktuellem zu neuem Spaltennamen, angewendet vor " + "der Typumwandlung. Nicht genannte Spalten behalten ihren Namen. " + "Die resultierenden Namen müssen eindeutig bleiben.", + zh="当前列名到新列名的映射，在类型转换前应用。未提及的列保留原名。" + "结果列名必须保持唯一。", + ), + alias=MultilingualString( + en="Column renames", + es="Renombres de columnas", + pt="Renomeações de colunas", + de="Spaltenumbenennungen", + zh="列重命名", + ), + ) # type: ignore + + +class ApplyDatasetSchemaUnit(BaseUnit): + """Rename the dataset's columns and cast them to their declared types. + + Takes the type declaration from the context (``inferred_types``) rather than + from its own configuration, so a single path feeds it whether the types were + inferred upstream or handed in by whoever set up the run. Renaming and + casting are one unit because a rename has to carry the declared types with + it: the type declared for ``Species`` has to end up on ``Variety``, not be + re-inferred from the data. + + ``validate`` rejects a declaration that names columns the dataset does not + have. Without that check the mismatch is silent — the underlying transform + passes unknown columns through untouched — so a stale declaration would + quietly leave columns with inferred types instead of the requested ones. + + Note the deliberate asymmetry: the unit does **not** republish a renamed + ``inferred_types``. After a rename that key describes column names that no + longer exist, and republishing it would invite a second consumer to trust a + declaration that no longer matches the dataset. Whoever needs types after + this unit reads them off the dataset. + """ + + SCHEMA = ApplyDatasetSchemaSchema + + REQUIRES = ("dataset", "inferred_types") + PROVIDES = ("dataset",) + + def validate(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + schema = ctx.require("inferred_types") + + unknown = sorted(set(schema) - set(dataset.column_names)) + if unknown: + raise JobError( + f"The declared types name columns the dataset does not have: {unknown}" + ) + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import ( + transform_dataset_with_schema, + ) + + dataset = ctx.require("dataset") + schema = ctx.require("inferred_types") + renames = self.config.get("column_renames") + + if renames: + dataset, schema = _rename_columns(dataset, schema, renames) + + ctx.put("dataset", transform_dataset_with_schema(dataset, schema)) + + +def _rename_columns(dataset, schema: dict, renames: dict): + """Return the dataset with renamed columns and the schema remapped to match. + + A plain helper: it takes and returns values and never touches the context, so + the unit's declared contract cannot hide inside it. + """ + original_names = dataset.arrow_table.schema.names + new_names = [renames.get(column, column) for column in original_names] + + if len(new_names) != len(set(new_names)): + duplicate_names = set() + seen = set() + for name in new_names: + if name in seen: + duplicate_names.add(name) + else: + seen.add(name) + raise JobError( + "Invalid column_renames: resulting column names contain duplicates: " + f"{sorted(duplicate_names)}" + ) + + arrow_table = dataset.arrow_table.rename_columns(new_names) + renamed = dataset.__class__( + arrow_table, + splits=dataset.splits, + types=dataset.types, + ) + remapped_schema = {renames.get(column, column): schema[column] for column in schema} + return renamed, remapped_schema diff --git a/DashAI/back/units/compute_dataset_metadata_unit.py b/DashAI/back/units/compute_dataset_metadata_unit.py new file mode 100644 index 000000000..d903b5263 --- /dev/null +++ b/DashAI/back/units/compute_dataset_metadata_unit.py @@ -0,0 +1,136 @@ +"""Unit that computes the metadata stored alongside a dataset.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema, bool_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 + +log = logging.getLogger(__name__) + +#: Metadata keys produced by ``compute_metadata`` on top of the base ones. +#: +#: They are expensive (correlations and per-column statistics over the whole +#: dataset) and optional, so the unit both skips computing them and strips any it +#: finds already present when they were not asked for. +EXTENDED_METADATA_KEYS = ( + "general_info", + "numeric_stats", + "categorical_stats", + "text_stats", + "quality_info", + "correlations", +) + + +class ComputeDatasetMetadataSchema(BaseSchema): + compute_metadata: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether to compute the extended exploratory metadata " + "(per-column statistics, quality report, correlations) on top of " + "the base one. Disabling it makes large datasets much faster to " + "store, at the cost of the exploration views.", + es="Si se debe calcular la metadata exploratoria extendida " + "(estadísticas por columna, informe de calidad, correlaciones) " + "además de la básica. Desactivarla hace mucho más rápido el " + "guardado de conjuntos de datos grandes, a costa de las vistas de " + "exploración.", + pt="Se deve calcular a metadata exploratória estendida " + "(estatísticas por coluna, relatório de qualidade, correlações) " + "além da básica. Desativá-la torna o armazenamento de grandes " + "conjuntos de dados muito mais rápido, ao custo das vistas de " + "exploração.", + de="Ob die erweiterten explorativen Metadaten (Spaltenstatistiken, " + "Qualitätsbericht, Korrelationen) zusätzlich zu den Basisdaten " + "berechnet werden. Das Deaktivieren beschleunigt das Speichern " + "großer Datensätze erheblich, auf Kosten der Explorationsansichten.", + zh="是否在基础元数据之外计算扩展的探索性元数据（各列统计、质量报告、" + "相关性）。禁用后可大幅加快大型数据集的存储速度，但会失去探索视图。", + ), + alias=MultilingualString( + en="Compute extended metadata", + es="Calcular metadata extendida", + pt="Calcular metadata estendida", + de="Erweiterte Metadaten berechnen", + zh="计算扩展元数据", + ), + ) # type: ignore + trust_inherited_metadata: schema_field( + bool_field(), + placeholder=False, + description=MultilingualString( + en="Whether the metadata the dataset already carries can be reused " + "as-is. Only safe when the data has not changed since that metadata " + "was computed; otherwise every value is recomputed from scratch.", + es="Si la metadata que el conjunto de datos ya trae puede reusarse " + "tal cual. Solo es seguro cuando los datos no cambiaron desde que " + "esa metadata se calculó; si no, todo se recalcula desde cero.", + pt="Se a metadata que o conjunto de dados já carrega pode ser " + "reutilizada como está. Só é seguro quando os dados não mudaram " + "desde que essa metadata foi calculada; caso contrário, tudo é " + "recalculado do zero.", + de="Ob die vom Datensatz mitgeführten Metadaten unverändert " + "weiterverwendet werden können. Nur sicher, wenn sich die Daten " + "seit deren Berechnung nicht geändert haben; andernfalls wird alles " + "neu berechnet.", + zh="数据集已携带的元数据是否可以原样重用。仅当数据自该元数据计算后" + "未发生变化时才安全；否则将全部重新计算。", + ), + alias=MultilingualString( + en="Trust existing metadata", + es="Confiar en la metadata existente", + pt="Confiar na metadata existente", + de="Vorhandene Metadaten vertrauen", + zh="信任现有元数据", + ), + ) # type: ignore + + +class ComputeDatasetMetadataUnit(BaseUnit): + """Fill in the metadata a dataset carries in its ``splits`` mapping. + + Two independent knobs, because "how much metadata" and "can the metadata + already there be trusted" are different questions: + + * ``compute_metadata`` picks the depth: base only (column names, row count, + NaN counts) or base plus the extended exploratory fields. + * ``trust_inherited_metadata`` says the dataset arrived carrying metadata + that still describes it — the case of a copy whose bytes never changed. + Then the unit fills only what is missing instead of recomputing. + + Either way the requested depth is what ends up stored: asking for base only + always strips the extended keys, even ones inherited from elsewhere, so the + result never depends on where the dataset came from. + + Re-publishes ``dataset`` although it mutates it in place: the key is what + makes the unit chainable, and the contract audit reads the ``ctx.put`` call. + """ + + SCHEMA = ComputeDatasetMetadataSchema + + REQUIRES = ("dataset",) + PROVIDES = ("dataset",) + + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + + compute_extended = self.config.get("compute_metadata", True) + trust_inherited = self.config.get("trust_inherited_metadata", False) + + if compute_extended: + has_extended = any(key in dataset.splits for key in EXTENDED_METADATA_KEYS) + if not (trust_inherited and has_extended): + dataset.compute_metadata() + else: + if not (trust_inherited and "total_rows" in dataset.splits): + dataset.compute_base_metadata() + # Strip extended keys the dataset may have arrived with, so the + # stored metadata matches what was asked for rather than the + # history of the file. + for stale_key in EXTENDED_METADATA_KEYS: + dataset.splits.pop(stale_key, None) + + ctx.put("dataset", dataset) diff --git a/DashAI/back/units/infer_dataset_types_unit.py b/DashAI/back/units/infer_dataset_types_unit.py new file mode 100644 index 000000000..58f3785c3 --- /dev/null +++ b/DashAI/back/units/infer_dataset_types_unit.py @@ -0,0 +1,76 @@ +"""Unit that works out what type each column of a dataset holds.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema, 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__) + + +class InferDatasetTypesSchema(BaseSchema): + method: schema_field( + string_field(), + placeholder="DashAIPtype", + description=MultilingualString( + en="Inference method used when the dataset does not already carry " + "types of its own.", + es="Método de inferencia que se usa cuando el conjunto de datos no " + "trae tipos propios.", + pt="Método de inferência usado quando o conjunto de dados não traz " + "tipos próprios.", + de="Inferenzmethode, die verwendet wird, wenn der Datensatz keine " + "eigenen Typen mitbringt.", + zh="当数据集本身未携带类型时使用的推断方法。", + ), + alias=MultilingualString( + en="Inference method", + es="Método de inferencia", + pt="Método de inferência", + de="Inferenzmethode", + zh="推断方法", + ), + ) # type: ignore + + +class InferDatasetTypesUnit(BaseUnit): + """Publish a type declaration for every column of the dataset. + + Prefers the types the dataset already carries: a dataloader that can read + them from the source (a Parquet schema, a typed database column) knows better + than any inference over the values, so re-inferring would throw that away. + Only when there are none does it fall back to inferring from the data. + + Publishes the declaration as a plain reference rather than applying it, so + the decision and the transformation stay separable — the same declaration can + be shown to a user for review before anything is cast. + + Note the reference names columns, so it goes stale the moment something + renames or drops one. It is meant to be consumed by the next step, not + carried across a chain of transformations. + """ + + SCHEMA = InferDatasetTypesSchema + + REQUIRES = ("dataset",) + PROVIDES = ("inferred_types",) + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.types.inf.type_inference import infer_types + + dataset = ctx.require("dataset") + + if dataset.types: + schema = { + column: dashai_type.to_string() + for column, dashai_type in dataset.types.items() + } + else: + schema = infer_types( + dataset.to_pandas(), + method=self.config.get("method", "DashAIPtype"), + ) + + ctx.put_ref("inferred_types", schema) diff --git a/DashAI/back/units/load_datafile_dataset_unit.py b/DashAI/back/units/load_datafile_dataset_unit.py new file mode 100644 index 000000000..2b97c4f28 --- /dev/null +++ b/DashAI/back/units/load_datafile_dataset_unit.py @@ -0,0 +1,178 @@ +"""Unit that reads a dataset out of an already downloaded datafile.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + none_type, + 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 + +log = logging.getLogger(__name__) + + +class LoadDatafileDatasetSchema(BaseSchema): + dataloader: schema_field( + component_field(parent="BaseDataLoader"), + placeholder={"component": "CSVDataLoader", "params": {"separator": ","}}, + description=MultilingualString( + en="Reader used to parse the downloaded file, together with its own " + "configuration.", + es="Lector usado para interpretar el archivo descargado, junto con " + "su propia configuración.", + pt="Leitor usado para interpretar o ficheiro descarregado, junto com " + "a sua própria configuração.", + de="Leser zum Parsen der heruntergeladenen Datei samt eigener " + "Konfiguration.", + zh="用于解析已下载文件的读取器及其自身配置。", + ), + alias=MultilingualString( + en="Data loader", + es="Cargador de datos", + pt="Carregador de dados", + de="Datenlader", + zh="数据加载器", + ), + ) # type: ignore + datafile_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the completed download to read from. It has to have " + "finished successfully; a download still running has no files yet.", + es="Identificador de la descarga completada desde la que leer. Tiene " + "que haber terminado con éxito; una descarga en curso todavía no " + "tiene archivos.", + pt="Identificador da descarga concluída de onde ler. Tem de ter " + "terminado com sucesso; uma descarga em curso ainda não tem " + "ficheiros.", + de="Kennung des abgeschlossenen Downloads, aus dem gelesen wird. Er " + "muss erfolgreich beendet sein; ein laufender Download hat noch " + "keine Dateien.", + zh="要读取的已完成下载的标识符。必须已成功完成；仍在进行的下载尚无文件。", + ), + alias=MultilingualString( + en="Downloaded file", + es="Archivo descargado", + pt="Ficheiro descarregado", + de="Heruntergeladene Datei", + zh="已下载文件", + ), + ) # type: ignore + selected_file: schema_field( + none_type(string_field()), + placeholder=None, + description=MultilingualString( + en="Which file inside the download to read, relative to its root. " + "Leave empty to take the first one, ignoring hidden files.", + es="Qué archivo dentro de la descarga leer, relativo a su raíz. " + "Dejar vacío para tomar el primero, ignorando archivos ocultos.", + pt="Qual ficheiro dentro da descarga ler, relativo à sua raiz. " + "Deixar vazio para tomar o primeiro, ignorando ficheiros ocultos.", + de="Welche Datei innerhalb des Downloads gelesen wird, relativ zu " + "dessen Wurzel. Leer lassen, um die erste zu nehmen; versteckte " + "Dateien werden übersprungen.", + zh="读取下载内容中的哪个文件（相对于其根目录）。留空则取第一个，忽略" + "隐藏文件。", + ), + alias=MultilingualString( + en="File", + es="Archivo", + pt="Ficheiro", + de="Datei", + zh="文件", + ), + ) # type: ignore + + +class LoadDatafileDatasetUnit(BaseUnit): + """Parse a dataset out of a download that already completed. + + Separate from ``LoadUploadedDatasetUnit`` even though both end in the same + reader call, because finding *what* to read is the whole job here: the + download is a directory tree recorded in the database, not a file the caller + hands over. The unit re-reads that row in its own read-only session, the same + way ``LoadDatasetUnit`` does, and never writes to it. + + Publishes only ``dataset``, for the same reason as its sibling: an imported + file is not a stored dataset yet, so there is no id to correlate and no path + to save back to. + """ + + SCHEMA = LoadDatafileDatasetSchema + + PROVIDES = ("dataset",) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.core.enums.status import DatafileStatus + from DashAI.back.dependencies.database.models import Datafile + + component_registry = di["component_registry"] + session_factory = di["session_factory"] + + datafile_id = self.config["datafile_id"] + selected_file = self.config.get("selected_file") + dataloader_config = self.config["dataloader"] + + with session_factory() as db: + datafile = db.get(Datafile, datafile_id) + # A missing row and an unfinished download are the same problem from + # here: there are no files to read either way. + if datafile is None or datafile.status != DatafileStatus.READY: + raise JobError(f"Datafile {datafile_id} is not ready.") + work_dir = datafile.local_path + + source = _resolve_source_file(work_dir, selected_file) + + dataloader_name = dataloader_config["component"] + registry = component_registry.registry.get("DataLoader", {}) + if dataloader_name not in registry: + raise JobError(f"DataLoader '{dataloader_name}' not found in registry.") + dataloader = registry[dataloader_name]["class"]() + + log.debug("Loading hub dataset from %s using %s", source, dataloader_name) + ctx.put( + "dataset", + dataloader.load_data( + filepath_or_buffer=source, + temp_path=work_dir, + params=dataloader_config.get("params") or {}, + n_sample=None, + ), + ) + + +def _resolve_source_file(work_dir: str, selected_file) -> str: + """Pick the file to read inside a completed download. + + A plain helper: takes and returns values, never touches the context. + + With no explicit choice it walks the tree and takes the first file in sorted + order, skipping anything under a dot-prefixed path component — download tools + leave metadata directories (``.cache``, ``.git``) behind that sort before the + real data and would otherwise win. + """ + from pathlib import Path + + if selected_file: + return str(Path(work_dir) / selected_file) + + base = Path(work_dir) + files = sorted( + str(path) + for path in base.rglob("*") + if path.is_file() + and not any(part.startswith(".") for part in path.relative_to(base).parts) + ) + if not files: + raise JobError("Hub download directory is empty.") + return files[0] diff --git a/DashAI/back/units/load_uploaded_dataset_unit.py b/DashAI/back/units/load_uploaded_dataset_unit.py new file mode 100644 index 000000000..0b708d0d2 --- /dev/null +++ b/DashAI/back/units/load_uploaded_dataset_unit.py @@ -0,0 +1,154 @@ +"""Unit that reads a dataset from an uploaded file or a URL.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + none_type, + 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__) + + +class LoadUploadedDatasetSchema(BaseSchema): + dataloader: schema_field( + component_field(parent="BaseDataLoader"), + placeholder={"component": "CSVDataLoader", "params": {"separator": ","}}, + description=MultilingualString( + en="Reader used to parse the source, together with its own " + "configuration (delimiter, sheet, encoding, and so on).", + es="Lector usado para interpretar el origen, junto con su propia " + "configuración (delimitador, hoja, codificación, etc.).", + pt="Leitor usado para interpretar a origem, junto com a sua própria " + "configuração (delimitador, folha, codificação, etc.).", + de="Leser zum Parsen der Quelle samt eigener Konfiguration " + "(Trennzeichen, Blatt, Kodierung usw.).", + zh="用于解析数据源的读取器及其自身配置（分隔符、工作表、编码等）。", + ), + alias=MultilingualString( + en="Data loader", + es="Cargador de datos", + pt="Carregador de dados", + de="Datenlader", + zh="数据加载器", + ), + ) # type: ignore + source: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Where to read from: a path to a local file or a URL. Archives " + "are downloaded and extracted before being parsed.", + es="Desde dónde leer: una ruta a un archivo local o una URL. Los " + "archivos comprimidos se descargan y extraen antes de interpretarse.", + pt="De onde ler: um caminho para um ficheiro local ou um URL. Os " + "arquivos comprimidos são descarregados e extraídos antes de serem " + "interpretados.", + de="Woher gelesen wird: ein Pfad zu einer lokalen Datei oder eine " + "URL. Archive werden vor dem Parsen heruntergeladen und entpackt.", + zh="读取来源：本地文件路径或 URL。压缩包会先下载并解压再解析。", + ), + alias=MultilingualString( + en="Source", + es="Origen", + pt="Origem", + de="Quelle", + zh="来源", + ), + ) # type: ignore + temp_path: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Scratch directory for downloads and extracted archives. Whoever " + "sets it up is responsible for removing it afterwards.", + es="Directorio temporal para descargas y archivos extraídos. Quien " + "lo crea es responsable de borrarlo después.", + pt="Diretório temporário para descargas e ficheiros extraídos. Quem " + "o cria é responsável por removê-lo depois.", + de="Arbeitsverzeichnis für Downloads und entpackte Archive. Wer es " + "anlegt, ist für das Entfernen verantwortlich.", + zh="用于下载和解压归档的临时目录。创建者负责事后清理。", + ), + alias=MultilingualString( + en="Temporary path", + es="Ruta temporal", + pt="Caminho temporário", + de="Temporärer Pfad", + zh="临时路径", + ), + ) # type: ignore + n_sample: schema_field( + none_type(int_field(gt=0)), + placeholder=None, + description=MultilingualString( + en="Read only this many rows instead of the whole source. Leave " + "empty to read everything.", + es="Leer solo esta cantidad de filas en vez de todo el origen. " + "Dejar vacío para leer todo.", + pt="Ler apenas esta quantidade de linhas em vez de toda a origem. " + "Deixar vazio para ler tudo.", + de="Nur so viele Zeilen lesen statt der gesamten Quelle. Leer " + "lassen, um alles zu lesen.", + zh="仅读取这么多行而非整个数据源。留空表示全部读取。", + ), + alias=MultilingualString( + en="Row sample", + es="Muestra de filas", + pt="Amostra de linhas", + de="Zeilenstichprobe", + zh="行采样", + ), + ) # type: ignore + + +class LoadUploadedDatasetUnit(BaseUnit): + """Parse a file or URL into a dataset with the chosen reader. + + The counterpart of ``LoadDatasetUnit``: that one materialises something DashAI + already stores, this one brings in data that is not a dataset yet. So it + publishes only ``dataset`` — there is no stored id to correlate against and no + path to save back to, because nothing decided yet where the result belongs. + + Types are not applied here. What the reader produces is whatever the source + suggests; declaring and casting types is a separate step, so the same load can + be reviewed before anything is committed to. + """ + + SCHEMA = LoadUploadedDatasetSchema + + PROVIDES = ("dataset",) + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + component_registry = di["component_registry"] + + dataloader_config = self.config["dataloader"] + source = self.config["source"] + + # Raises KeyError with the registry's own wording when the reader does + # not exist. That message reaches the user, so it is not reworded here. + dataloader = component_registry[dataloader_config["component"]]["class"]() + + log.debug( + "Loading dataset from %s using %s", + source, + dataloader_config["component"], + ) + ctx.put( + "dataset", + dataloader.load_data( + filepath_or_buffer=source, + temp_path=self.config["temp_path"], + params=dataloader_config.get("params") or {}, + n_sample=self.config.get("n_sample"), + ), + ) diff --git a/DashAI/back/units/save_dataset_to_path_unit.py b/DashAI/back/units/save_dataset_to_path_unit.py new file mode 100644 index 000000000..6c5c80649 --- /dev/null +++ b/DashAI/back/units/save_dataset_to_path_unit.py @@ -0,0 +1,78 @@ +"""Unit that writes the dataset in the context to a destination of its own.""" + +import logging + +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 + +log = logging.getLogger(__name__) + + +class SaveDatasetToPathSchema(BaseSchema): + path: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Destination directory for the dataset. Unlike saving in place, " + "this writes wherever it is told, so it can store a dataset that " + "was loaded from somewhere else entirely.", + es="Directorio de destino del conjunto de datos. A diferencia de " + "guardar en su lugar, escribe donde se le indique, así que puede " + "almacenar un conjunto de datos cargado desde otro origen.", + pt="Diretório de destino do conjunto de dados. Ao contrário de " + "guardar no lugar, escreve onde lhe for indicado, pelo que pode " + "armazenar um conjunto de dados carregado de outra origem.", + de="Zielverzeichnis für den Datensatz. Anders als beim Speichern am " + "Ursprungsort schreibt diese Einheit dorthin, wo es ihr gesagt " + "wird, und kann so einen anderswo geladenen Datensatz ablegen.", + zh="数据集的目标目录。与原地保存不同，它会写入指定位置，因此可以存储" + "从其他来源加载的数据集。", + ), + alias=MultilingualString( + en="Destination path", + es="Ruta de destino", + pt="Caminho de destino", + de="Zielpfad", + zh="目标路径", + ), + ) # type: ignore + + +class SaveDatasetToPathUnit(BaseUnit): + """Write the dataset to a destination given in the configuration. + + The sibling of ``SaveDatasetUnit``, and deliberately a separate unit rather + than a flag on it. ``SaveDatasetUnit`` saves back to ``dataset_path`` — where + the load came from — which is what makes editing a dataset in place safe. This + one stores the dataset as something new, so it never reads ``dataset_path``: + if it did, a flow that loads a working copy and registers the result as a + fresh dataset would overwrite the copy it read. + + Declares no outputs — its result is on disk, not in the context. + + The error message keeps the original exception's text. A failure here is an + infrastructure failure (out of space, permissions, a path the filesystem + rejects), and only the message of the outermost error reaches the user: the + job queue stores ``str(exc)``, never the ``__cause__`` chain. Swallowing it + would drop the diagnosis exactly when it is needed. + """ + + SCHEMA = SaveDatasetToPathSchema + + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import save_dataset + + dataset = ctx.require("dataset") + path = self.config["path"] + + try: + save_dataset(dataset, path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not save dataset to path {path}: {e}") from e diff --git a/DashAI/back/units/save_dataset_unit.py b/DashAI/back/units/save_dataset_unit.py index 50d9c9bc7..519482edc 100644 --- a/DashAI/back/units/save_dataset_unit.py +++ b/DashAI/back/units/save_dataset_unit.py @@ -16,6 +16,12 @@ class SaveDatasetUnit(BaseUnit): whichever unit loaded the dataset, so a save can never land somewhere the load did not come from. Declares no outputs — its result is on disk, not in the context. + + The error message keeps the original exception's text. A failure here is an + infrastructure failure (out of space, permissions, a path the filesystem + rejects), and only the message of the outermost error reaches the user: the + job queue stores ``str(exc)``, never the ``__cause__`` chain. Swallowing it + would drop the diagnosis exactly when it is needed. """ REQUIRES = ("dataset", "dataset_path") @@ -31,4 +37,4 @@ def execute(self, ctx: ExecutionContext) -> None: save_dataset(dataset, dataset_path) except Exception as e: log.exception(e) - raise JobError(f"Can not save dataset to path {dataset_path}") from e + raise JobError(f"Can not save dataset to path {dataset_path}: {e}") from e diff --git a/tests/back/api/test_dataset_job.py b/tests/back/api/test_dataset_job.py new file mode 100644 index 000000000..f20e7caa7 --- /dev/null +++ b/tests/back/api/test_dataset_job.py @@ -0,0 +1,1094 @@ +"""End-to-end regression net for ``DatasetJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so the refactor has something to be measured against. +The assertions are deliberately explicit — exact status values, exact row/column +counts, exact keys in ``splits.json``, exact error message text — 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. + +``DatasetJob`` has the widest blast radius of any job in the repo: it is the +bootstrap fixture of half the suite (``tests/back/api/conftest.py``, +``tests/back/models/conftest.py``, ``tests/back/api/test_predict_api.py``, +``tests/back/types/load_preview_test.py``). Everything here has to keep passing +through every slice of the refactor. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import os +import shutil +from pathlib import Path +from typing import Optional + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import DatafileStatus, DatasetStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import Datafile, Dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.job.dataset_job import DatasetJob + +IRIS_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Species", +] +IRIS_ROWS = 150 + +IRIS_SCHEMA = { + "SepalLengthCm": {"type": "Float", "dtype": "float64"}, + "SepalWidthCm": {"type": "Float", "dtype": "float64"}, + "PetalLengthCm": {"type": "Float", "dtype": "float64"}, + "PetalWidthCm": {"type": "Float", "dtype": "float64"}, + "Species": {"type": "Categorical", "dtype": "string"}, +} + +#: The extended EDA keys ``compute_metadata=False`` has to leave out. +EXTENDED_KEYS = ( + "general_info", + "numeric_stats", + "categorical_stats", + "text_stats", + "quality_info", + "correlations", +) + +IRIS_CSV = Path(__file__).parent / "iris.csv" + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # + + +def _session_factory(client): + return client.app.container["session_factory"] + + +def _new_dataset_row(client, name: str) -> int: + """Create a ``Dataset`` row in DELIVERED, the state the job expects.""" + with _session_factory(client)() as db: + entry = Dataset(name=name, file_path="") + entry.set_status_as_delivered() + db.add(entry) + db.commit() + db.refresh(entry) + return entry.id + + +def _stored(client, dataset_id: int) -> Optional[dict]: + """Read the ``Dataset`` row straight from the database. + + No ``start_time`` / ``end_time``: unlike ``Run`` or ``Explorer``, this table + declares no such columns. The status setters used to assign them anyway, + which only set attributes that were dropped at commit. + """ + with _session_factory(client)() as db: + row = db.get(Dataset, dataset_id) + if row is None: + return None + return { + "status": row.status, + "file_path": row.file_path, + "total_rows": row.total_rows, + "total_columns": row.total_columns, + "last_modified": row.last_modified, + } + + +def _run_job(client, dataset_id: int, params: dict, **extra) -> None: + """Run ``DatasetJob`` synchronously, the way every fixture in the suite does.""" + kwargs = { + "dataset_id": dataset_id, + "url": "", + "params": params, + "file_path": IRIS_CSV, + **extra, + } + DatasetJob(job_type="DatasetJob", kwargs=kwargs).run() + + +def _csv_params(name: str, **overrides) -> dict: + params = { + "dataloader": "CSVDataLoader", + "separator": ",", + "name": name, + "schema": IRIS_SCHEMA, + } + params.update(overrides) + return params + + +def _splits(file_path: str) -> dict: + with open(Path(file_path) / "dataset" / "splits.json", encoding="utf-8") as f: + return json.load(f) + + +def _cleanup(client, dataset_id: int) -> None: + row = _stored(client, dataset_id) + if row and row["file_path"]: + shutil.rmtree(row["file_path"], ignore_errors=True) + + +# --------------------------------------------------------------------------- # +# status bookkeeping +# --------------------------------------------------------------------------- # + + +class TestTheRowIsNeverLeftStuckInStarted: + """Whatever goes wrong, the row must not stay in STARTED. + + Nothing else moves it: Huey's ``SIGNAL_ERROR`` writes to its own + ``task_copy`` table and never to the ``Dataset`` row + (``huey_job_queue.py:201-211``), so an exception the job's own handler does + not see leaves the dataset advertised as in-progress forever and the UI + spinning on work that already died. + + Both cases below used to escape — the handler only caught ``JobError``, and + the final database block only caught ``SQLAlchemyError``. They are the + reason the handler now catches ``Exception``. + """ + + def test_a_mkdir_failure_that_is_not_file_exists_still_marks_the_row( + self, client, monkeypatch + ): + """Creating the destination folder can fail with more than + ``FileExistsError``: no permission, no space, a path the filesystem + rejects. The status was already committed as STARTED by then.""" + dataset_id = _new_dataset_row(client, "stuck_on_mkdir") + + def deny(self, *args, **kwargs): + raise PermissionError("cannot create the dataset folder") + + monkeypatch.setattr(Path, "mkdir", deny) + + with pytest.raises(PermissionError): + _run_job(client, dataset_id, _csv_params("stuck_on_mkdir")) + + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + def test_a_row_deleted_mid_run_is_reported_instead_of_crashing(self, client): + """The row is deletable through the API while the job runs, so the final + block can find it gone. It used to assign to ``None`` and raise + ``AttributeError`` from inside a block that only caught database errors. + """ + dataset_id = _new_dataset_row(client, "stuck_on_deleted_row") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def save_then_delete(dataset, path): + real_save(dataset, path) + with _session_factory(client)() as db: + db.delete(db.get(Dataset, dataset_id)) + db.commit() + + dashai_dataset.save_dataset = save_then_delete + try: + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("stuck_on_deleted_row")) + finally: + dashai_dataset.save_dataset = real_save + + assert str(excinfo.value) == (f"Dataset with ID {dataset_id} no longer exists.") + # The row is gone, so there is no status to write — but the data written + # to disk has to go with it instead of being orphaned. + assert _stored(client, dataset_id) is None + assert _folders_in(datasets_dir) == before + + def test_a_non_database_failure_in_the_final_block_still_marks_the_row( + self, client + ): + """The mirror of the case above for a row that does still exist: the + block catches ``SQLAlchemyError`` by name, so anything else has to be + handled further out rather than escaping.""" + dataset_id = _new_dataset_row(client, "stuck_on_final_block") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + real_set_finished = Dataset.set_status_as_finished + + def boom(self): + raise RuntimeError("not a database error") + + Dataset.set_status_as_finished = boom + try: + with pytest.raises(RuntimeError): + _run_job(client, dataset_id, _csv_params("stuck_on_final_block")) + finally: + Dataset.set_status_as_finished = real_set_finished + + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + assert _folders_in(datasets_dir) == before + + +# --------------------------------------------------------------------------- # +# file / URL branch — the happy path +# --------------------------------------------------------------------------- # + + +def test_the_file_branch_finishes_and_writes_every_column_of_the_row(client): + """Status transitions, the row's columns, and the dataset on disk.""" + dataset_id = _new_dataset_row(client, "net_file_happy") + before = _stored(client, dataset_id) + assert before["status"] == DatasetStatus.DELIVERED + + _run_job(client, dataset_id, _csv_params("net_file_happy")) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + assert row["last_modified"] > before["last_modified"] + + # The path is a realpath under DATASETS_PATH, and the dataset is really there. + datasets_path = client.app.container["config"]["DATASETS_PATH"] + assert Path(row["file_path"]).parent == Path(os.path.realpath(datasets_path)) + + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_declared_schema_lands_on_the_stored_types(client): + """``transform_dataset_with_schema`` is what makes Species categorical.""" + dataset_id = _new_dataset_row(client, "net_file_types") + + _run_job(client, dataset_id, _csv_params("net_file_types")) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + types = dataset.types + + assert type(types["Species"]).__name__ == "Categorical" + for column in IRIS_COLUMNS[:4]: + assert type(types[column]).__name__ == "Float" + + _cleanup(client, dataset_id) + + +def test_inferred_types_in_the_params_win_over_inference(client): + """``params["inferred_types"]`` is the first branch of the schema cascade.""" + dataset_id = _new_dataset_row(client, "net_file_inferred") + # Species declared as plain text instead of categorical: only an honoured + # override can produce this. + override = dict(IRIS_SCHEMA, Species={"type": "Text", "dtype": "string"}) + + _run_job( + client, + dataset_id, + _csv_params("net_file_inferred", inferred_types=override), + ) + + row = _stored(client, dataset_id) + types = load_dataset(f"{row['file_path']}/dataset").types + assert type(types["Species"]).__name__ != "Categorical" + + _cleanup(client, dataset_id) + + +def test_the_schema_is_inferred_when_the_params_declare_nothing(client): + """Third branch of the cascade: ``infer_types(..., "DashAIPtype")``.""" + dataset_id = _new_dataset_row(client, "net_file_no_schema") + + _run_job( + client, + dataset_id, + {"dataloader": "CSVDataLoader", "separator": ",", "name": "net_file_no_schema"}, + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + # Inference ran and produced a type for every column. + assert set(dataset.types) == set(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +def test_column_renames_rewrite_the_columns_and_remap_the_schema(client): + """The rename remaps ``schema`` too (lines 291-297), so the renamed column + keeps the type its old name declared.""" + dataset_id = _new_dataset_row(client, "net_file_renames") + + _run_job( + client, + dataset_id, + _csv_params( + "net_file_renames", + inferred_types=IRIS_SCHEMA, + column_renames={"Species": "Variety", "SepalLengthCm": "SepalLength"}, + ), + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + + assert dataset.column_names == [ + "SepalLength", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", + "Variety", + ] + # The type followed the rename rather than being re-inferred. + assert type(dataset.types["Variety"]).__name__ == "Categorical" + assert row["total_columns"] == len(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# metadata policy +# --------------------------------------------------------------------------- # + + +def test_compute_metadata_true_writes_base_and_extended_metadata(client): + dataset_id = _new_dataset_row(client, "net_meta_full") + + _run_job(client, dataset_id, _csv_params("net_meta_full", compute_metadata=True)) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + assert splits["column_names"] == IRIS_COLUMNS + assert "nan" in splits + for key in EXTENDED_KEYS: + assert key in splits, f"missing {key} when compute_metadata=True" + + _cleanup(client, dataset_id) + + +def test_compute_metadata_false_writes_base_metadata_only(client): + dataset_id = _new_dataset_row(client, "net_meta_base") + + _run_job(client, dataset_id, _csv_params("net_meta_base", compute_metadata=False)) + + row = _stored(client, dataset_id) + splits = _splits(row["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + assert splits["column_names"] == IRIS_COLUMNS + assert "nan" in splits + for key in EXTENDED_KEYS: + assert key not in splits, f"unexpected {key} when compute_metadata=False" + + # The row's columns come from splits, so they must survive the base-only path. + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + + _cleanup(client, dataset_id) + + +def test_omitting_the_flag_defaults_to_full_metadata(client): + """Backward compatibility: the 4 bootstrap fixtures never pass the flag.""" + dataset_id = _new_dataset_row(client, "net_meta_default") + + _run_job(client, dataset_id, _csv_params("net_meta_default")) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + for key in EXTENDED_KEYS: + assert key in splits + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# notebook branch +# --------------------------------------------------------------------------- # + + +@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 the notebook branch of the + job reads a private working copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "dataset job net"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _add_converter(client, notebook_id: int) -> int: + """Register a converter against the notebook. + + The job only asks *whether any Converter row exists* for the notebook + (lines 174-180); it never runs it. So a row is enough to flip + ``from_notebook_no_converters`` to False. + """ + response = client.post( + "/api/v1/converter/", + json={ + "notebook_id": notebook_id, + "converter": "StandardScaler", + "parameters": { + "order": 0, + "params": {}, + "scope": {"columns": [], "rows": []}, + "target": None, + }, + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def test_the_notebook_branch_copies_the_working_copy_into_a_new_dataset( + client, notebook +): + dataset_id = _new_dataset_row(client, "net_notebook_happy") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_happy"}, + notebook_id=notebook["id"], + file_path=None, + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + + # A new folder, not the notebook's own copy. + assert Path(row["file_path"]) != Path(notebook["file_path"]) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + assert len(dataset) == IRIS_ROWS + + # And the notebook is untouched. + assert load_dataset(f"{notebook['file_path']}/dataset").column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +def test_a_notebook_without_converters_reuses_the_source_metadata(client, notebook): + """No converters means the bytes match the source, so ``splits.json`` is + inherited rather than recomputed (lines 313-328).""" + dataset_id = _new_dataset_row(client, "net_notebook_inherit") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_inherit", "compute_metadata": True}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + # dataset_1 was built with the default (full metadata), so the inherited + # splits already carry the extended keys. + for key in EXTENDED_KEYS: + assert key in splits + assert splits["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_a_notebook_without_converters_still_drops_extended_when_asked( + client, notebook +): + """``compute_metadata=False`` purges the extended keys the source carried.""" + dataset_id = _new_dataset_row(client, "net_notebook_purge") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_purge", "compute_metadata": False}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + assert splits["total_rows"] == IRIS_ROWS + for key in EXTENDED_KEYS: + assert key not in splits, f"inherited {key} was not purged" + + _cleanup(client, dataset_id) + + +def test_a_notebook_with_converters_recomputes_the_metadata(client, notebook): + """One Converter row is enough to stop trusting the source's splits.""" + _add_converter(client, notebook["id"]) + dataset_id = _new_dataset_row(client, "net_notebook_recompute") + + _run_job( + client, + dataset_id, + {"name": "net_notebook_recompute", "compute_metadata": True}, + notebook_id=notebook["id"], + file_path=None, + ) + + splits = _splits(_stored(client, dataset_id)["file_path"]) + for key in EXTENDED_KEYS: + assert key in splits + assert splits["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_notebook_branch_never_applies_a_schema_or_renames(client, notebook): + """Lines 260-299 sit in the ``else`` of the notebook check, so neither + ``inferred_types`` nor ``column_renames`` has any effect here.""" + dataset_id = _new_dataset_row(client, "net_notebook_no_schema") + + _run_job( + client, + dataset_id, + { + "name": "net_notebook_no_schema", + "column_renames": {"Species": "ShouldBeIgnored"}, + }, + notebook_id=notebook["id"], + file_path=None, + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# hub branch +# --------------------------------------------------------------------------- # + + +@pytest.fixture(name="datafile") +def create_datafile(client, tmp_path_factory, request): + """A READY ``Datafile`` row pointing at a directory holding iris.csv. + + Built by hand: the hub branch only reads the row and the files on disk, so + no network is involved. ``dataset_id`` carries the test name because the + table has a ``UNIQUE(source_name, dataset_id)`` constraint and each test + gets its own row. + """ + work_dir = tmp_path_factory.mktemp("hub_download") + shutil.copy(IRIS_CSV, work_dir / "iris.csv") + + with _session_factory(client)() as db: + row = Datafile( + source_name="HuggingFaceDatasetSource", + dataset_id=f"net/iris/{request.node.name}", + name="net iris", + local_path=str(work_dir), + status=DatafileStatus.READY, + ) + db.add(row) + db.commit() + db.refresh(row) + return {"id": row.id, "local_path": str(work_dir)} + + +def _hub_params(name: str, datafile_id: int, **overrides) -> dict: + params = { + "name": name, + "dataloader": "CSVDataLoader", + "dataloader_params": {"separator": ","}, + "datafile_id": datafile_id, + } + params.update(overrides) + return params + + +def _run_hub_job(client, dataset_id: int, params: dict) -> None: + DatasetJob( + job_type="DatasetJob", + kwargs={ + "dataset_id": dataset_id, + "source_name": "HuggingFaceDatasetSource", + "dataset_source_id": "net/iris", + "params": params, + }, + ).run() + + +def test_the_hub_branch_finishes_from_an_explicit_selected_file(client, datafile): + dataset_id = _new_dataset_row(client, "net_hub_selected") + + _run_hub_job( + client, + dataset_id, + _hub_params("net_hub_selected", datafile["id"], selected_file="iris.csv"), + ) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + assert row["total_columns"] == len(IRIS_COLUMNS) + assert load_dataset(f"{row['file_path']}/dataset").column_names == IRIS_COLUMNS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_picks_the_first_file_when_none_is_selected(client, datafile): + """No ``selected_file``: the first entry of a sorted ``rglob`` wins.""" + dataset_id = _new_dataset_row(client, "net_hub_first") + + _run_hub_job(client, dataset_id, _hub_params("net_hub_first", datafile["id"])) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_ignores_dotted_files_when_picking(client, datafile): + """Any path component starting with a dot is skipped (lines 213-217). + + ``.hidden`` sorts before ``iris.csv``, so without the filter it would win + and the job would fail on an unreadable file. + """ + hidden = Path(datafile["local_path"]) / ".hidden" + hidden.mkdir() + (hidden / "junk.csv").write_text("not,a,dataset\n", encoding="utf-8") + + dataset_id = _new_dataset_row(client, "net_hub_hidden") + + _run_hub_job(client, dataset_id, _hub_params("net_hub_hidden", datafile["id"])) + + row = _stored(client, dataset_id) + assert row["status"] == DatasetStatus.FINISHED + assert row["total_rows"] == IRIS_ROWS + + _cleanup(client, dataset_id) + + +def test_the_hub_branch_applies_the_schema_and_the_renames(client, datafile): + """Unlike the notebook branch, hub imports do go through lines 260-299.""" + dataset_id = _new_dataset_row(client, "net_hub_renames") + + _run_hub_job( + client, + dataset_id, + _hub_params( + "net_hub_renames", + datafile["id"], + selected_file="iris.csv", + inferred_types=IRIS_SCHEMA, + column_renames={"Species": "Variety"}, + ), + ) + + row = _stored(client, dataset_id) + dataset = load_dataset(f"{row['file_path']}/dataset") + assert dataset.column_names[-1] == "Variety" + assert type(dataset.types["Variety"]).__name__ == "Categorical" + + _cleanup(client, dataset_id) + + +# --------------------------------------------------------------------------- # +# error messages, by exact text +# --------------------------------------------------------------------------- # +# +# These travel verbatim to ``task_copy.error_msg`` (``huey_job_queue.py:210`` +# stores ``str(exc)``) and from there to ``GET /job/status/{id}`` and the UI, so +# the text *is* the contract. Everything raised inside the try at line 153 comes +# back wrapped as ``Error loading dataset: ``. + + +def _datasets_dir(client) -> Path: + return Path(client.app.container["config"]["DATASETS_PATH"]) + + +def _folders_in(path: Path) -> set: + return {p.name for p in path.iterdir() if p.is_dir()} + + +def test_a_missing_dataset_row_is_reported_by_id(client): + """Raised before the status is touched, so nothing is left behind.""" + with pytest.raises(JobError) as excinfo: + _run_job(client, 987654, _csv_params("net_missing_row")) + + assert str(excinfo.value) == "Dataset with ID 987654 not found." + + +def test_a_file_exists_collision_names_the_generated_folder(client, monkeypatch): + dataset_id = _new_dataset_row(client, "net_collision") + + def collide(self, *args, **kwargs): + raise FileExistsError(self) + + monkeypatch.setattr(Path, "mkdir", collide) + + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("net_collision")) + + message = str(excinfo.value) + assert message.startswith("A dataset with the name ") + assert message.endswith(" already exists.") + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_missing_notebook_is_reported_with_the_original_wording(client): + """The wording says "has no associated dataset" even though what is missing + is the Notebook row itself. Preserved verbatim: it reaches the UI.""" + dataset_id = _new_dataset_row(client, "net_missing_notebook") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_missing_notebook"}, + notebook_id=987654, + file_path=None, + ) + + assert str(excinfo.value) == ( + "Error loading dataset: Notebook with ID 987654 has no associated dataset." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unreadable_notebook_copy_names_the_path_it_could_not_read(client, notebook): + """The notebook branch reads a stored dataset off disk. When that read fails + the message names the path instead of quoting the reader's exception. + + This is the one intentional wording change of the refactor: the branch now + goes through the same loading step as every other flow that materialises a + stored dataset, and that step reports failures by path. The path is the + actionable part — the reader's own text ("the arrow file is truncated") says + nothing about *which* file. + """ + dataset_id = _new_dataset_row(client, "net_unreadable_notebook") + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_load = dashai_dataset.load_dataset + + def fail_to_load(path): + raise OSError("the arrow file is truncated") + + dashai_dataset.load_dataset = fail_to_load + try: + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_unreadable_notebook"}, + notebook_id=notebook["id"], + file_path=None, + ) + finally: + dashai_dataset.load_dataset = real_load + + assert str(excinfo.value) == ( + f"Error loading dataset: Can not load dataset from path {notebook['file_path']}" + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unknown_dataloader_in_the_file_branch_reports_the_registry_key_error( + client, +): + """The file branch uses ``component_registry[name]``, whose ``KeyError`` + stringifies *with* the quotes.""" + dataset_id = _new_dataset_row(client, "net_bad_loader_file") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + _csv_params("net_bad_loader_file", dataloader="NoSuchDataLoader"), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: \"Component 'NoSuchDataLoader' does not exists " + 'in the registry."' + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_an_unknown_dataloader_in_the_hub_branch_has_its_own_wording(client, datafile): + """The hub branch looks the loader up itself, with a different message than + the file branch. Both have to survive the refactor.""" + dataset_id = _new_dataset_row(client, "net_bad_loader_hub") + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, + dataset_id, + _hub_params( + "net_bad_loader_hub", datafile["id"], dataloader="NoSuchDataLoader" + ), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: DataLoader 'NoSuchDataLoader' not found in registry." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_hub_import_without_a_datafile_id_says_so(client): + dataset_id = _new_dataset_row(client, "net_hub_no_id") + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, + dataset_id, + {"name": "net_hub_no_id", "dataloader": "CSVDataLoader"}, + ) + + assert str(excinfo.value) == ( + "Error loading dataset: datafile_id is required for hub imports." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_datafile_that_is_not_ready_is_reported_by_id(client, datafile): + dataset_id = _new_dataset_row(client, "net_hub_not_ready") + with _session_factory(client)() as db: + row = db.get(Datafile, datafile["id"]) + row.status = DatafileStatus.DOWNLOADING + db.commit() + + with pytest.raises(JobError) as excinfo: + _run_hub_job( + client, dataset_id, _hub_params("net_hub_not_ready", datafile["id"]) + ) + + assert str(excinfo.value) == ( + f"Error loading dataset: Datafile {datafile['id']} is not ready." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_missing_datafile_row_is_also_reported_as_not_ready(client): + """``hub_row is None`` and ``status != READY`` share one branch and one + message, so a nonexistent id reads as "not ready".""" + dataset_id = _new_dataset_row(client, "net_hub_missing_row") + + with pytest.raises(JobError) as excinfo: + _run_hub_job(client, dataset_id, _hub_params("net_hub_missing_row", 987654)) + + assert str(excinfo.value) == "Error loading dataset: Datafile 987654 is not ready." + + +def test_an_empty_hub_download_directory_says_so(client, tmp_path_factory): + dataset_id = _new_dataset_row(client, "net_hub_empty") + empty_dir = tmp_path_factory.mktemp("hub_empty") + with _session_factory(client)() as db: + row = Datafile( + source_name="HuggingFaceDatasetSource", + dataset_id="net/iris/empty", + name="empty", + local_path=str(empty_dir), + status=DatafileStatus.READY, + ) + db.add(row) + db.commit() + db.refresh(row) + datafile_id = row.id + + with pytest.raises(JobError) as excinfo: + _run_hub_job(client, dataset_id, _hub_params("net_hub_empty", datafile_id)) + + assert str(excinfo.value) == ( + "Error loading dataset: Hub download directory is empty." + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_column_renames_that_collide_list_the_duplicates_sorted(client): + dataset_id = _new_dataset_row(client, "net_dup_renames") + + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + _csv_params( + "net_dup_renames", + inferred_types=IRIS_SCHEMA, + column_renames={ + "SepalWidthCm": "Same", + "PetalWidthCm": "Same", + "Species": "SepalLengthCm", + }, + ), + ) + + assert str(excinfo.value) == ( + "Error loading dataset: Invalid column_renames: resulting column names " + "contain duplicates: ['Same', 'SepalLengthCm']" + ) + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_a_database_error_on_the_final_commit_is_reported_generically(client): + """The last block (lines 351-368) catches only ``SQLAlchemyError`` and + replaces it with a fixed message, dropping the database's own text. + + Note the asymmetry with everything above: this one is raised *outside* the + try at line 153, so it is not prefixed with "Error loading dataset: ". + """ + from sqlalchemy import exc as sa_exc + + dataset_id = _new_dataset_row(client, "net_db_error") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + real_set_finished = Dataset.set_status_as_finished + + def boom(self): + raise sa_exc.InvalidRequestError("the database said no") + + Dataset.set_status_as_finished = boom + try: + with pytest.raises(JobError) as excinfo: + _run_job(client, dataset_id, _csv_params("net_db_error")) + finally: + Dataset.set_status_as_finished = real_set_finished + + assert str(excinfo.value) == "Internal database error" + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + # This branch cleans up the folder too. + assert _folders_in(datasets_dir) == before + + +# --------------------------------------------------------------------------- # +# side effects of a failure +# --------------------------------------------------------------------------- # + + +def test_a_failure_after_the_folder_was_created_removes_it(client): + """The generated folder is created before loading, and the ``except`` at + line 345 has to clean it up — otherwise every failed import leaks a folder + under DATASETS_PATH.""" + dataset_id = _new_dataset_row(client, "net_leak_check") + datasets_dir = _datasets_dir(client) + before = _folders_in(datasets_dir) + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_leak_check", dataloader="NoSuchDataLoader"), + ) + + assert _folders_in(datasets_dir) == before + assert _stored(client, dataset_id)["file_path"] == "" + + +def test_a_failure_does_not_delete_a_folder_the_job_did_not_create(client, tmp_path): + """Re-importing into an existing dataset reuses that dataset's own folder + instead of creating one (the ``n_sample`` branch). + + The cleanup on failure removes the destination folder, which is right when + the job created it and destructive when it did not: the row survives the + failure pointing at a path whose contents would be gone. + """ + existing = tmp_path / "already_stored" + existing.mkdir() + (existing / "dataset").mkdir() + (existing / "dataset" / "data.arrow").write_text("payload", encoding="utf-8") + + with _session_factory(client)() as db: + entry = Dataset(name="net_reused_folder", file_path=str(existing)) + entry.set_status_as_delivered() + db.add(entry) + db.commit() + db.refresh(entry) + dataset_id = entry.id + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_reused_folder", dataloader="NoSuchDataLoader"), + n_sample=10, + ) + + assert (existing / "dataset" / "data.arrow").read_text( + encoding="utf-8" + ) == "payload" + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR + + +def test_the_temp_dir_is_removed_even_on_the_happy_path(client, tmp_path_factory): + """The ``finally`` removes ``temp_dir`` whether the job worked or not.""" + dataset_id = _new_dataset_row(client, "net_temp_happy") + temp_dir = tmp_path_factory.mktemp("net_temp_happy") + + _run_job(client, dataset_id, _csv_params("net_temp_happy"), temp_dir=str(temp_dir)) + + assert not temp_dir.exists() + assert _stored(client, dataset_id)["status"] == DatasetStatus.FINISHED + + _cleanup(client, dataset_id) + + +def test_the_temp_dir_is_removed_after_a_failure(client, tmp_path_factory): + dataset_id = _new_dataset_row(client, "net_temp_failure") + temp_dir = tmp_path_factory.mktemp("net_temp_failure") + + with pytest.raises(JobError): + _run_job( + client, + dataset_id, + _csv_params("net_temp_failure", dataloader="NoSuchDataLoader"), + temp_dir=str(temp_dir), + ) + + assert not temp_dir.exists() + + +def test_a_failed_notebook_import_leaves_the_notebook_intact(client, notebook): + """The cleanup ``rmtree`` targets the *generated* folder. If it ever pointed + at the loaded path instead, this is what would catch it: the notebook's own + working copy would be gone.""" + dataset_id = _new_dataset_row(client, "net_notebook_failure") + notebook_dataset = Path(notebook["file_path"]) / "dataset" + assert notebook_dataset.exists() + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def fail_to_save(dataset, path): + raise OSError("disk full") + + dashai_dataset.save_dataset = fail_to_save + try: + with pytest.raises(JobError) as excinfo: + _run_job( + client, + dataset_id, + {"name": "net_notebook_failure"}, + notebook_id=notebook["id"], + file_path=None, + ) + finally: + dashai_dataset.save_dataset = real_save + + # The message gained the destination path, and deliberately kept the + # original exception's text: a save failure is an infrastructure failure and + # only the outermost message reaches the user (``huey_job_queue.py:210`` + # stores ``str(exc)``, never the ``__cause__`` chain), so swallowing it would + # drop the diagnosis exactly when it matters. + message = str(excinfo.value) + assert message.startswith("Error loading dataset: Can not save dataset to path ") + assert message.endswith(": disk full") + + assert notebook_dataset.exists() + assert load_dataset(str(notebook_dataset)).column_names == IRIS_COLUMNS + assert _stored(client, dataset_id)["status"] == DatasetStatus.ERROR diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index 6443545d1..dd5726b8d 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -27,6 +27,12 @@ "PrepareExplanationDataUnit", "GenerateGlobalExplanationUnit", "GenerateLocalExplanationUnit", + "LoadUploadedDatasetUnit", + "LoadDatafileDatasetUnit", + "InferDatasetTypesUnit", + "ApplyDatasetSchemaUnit", + "ComputeDatasetMetadataUnit", + "SaveDatasetToPathUnit", } @@ -108,6 +114,30 @@ def test_unit_schemas_describe_their_configuration(units): assert set(units["GenerateGlobalExplanationUnit"]["schema"]["properties"]) == { "explainer_id" } + assert set(units["LoadUploadedDatasetUnit"]["schema"]["properties"]) == { + "dataloader", + "source", + "temp_path", + "n_sample", + } + assert set(units["LoadDatafileDatasetUnit"]["schema"]["properties"]) == { + "dataloader", + "datafile_id", + "selected_file", + } + assert set(units["InferDatasetTypesUnit"]["schema"]["properties"]) == {"method"} + # The type declaration arrives through the context, not the configuration, + # so the only thing to configure here is the renaming. + assert set(units["ApplyDatasetSchemaUnit"]["schema"]["properties"]) == { + "column_renames" + } + assert set(units["ComputeDatasetMetadataUnit"]["schema"]["properties"]) == { + "compute_metadata", + "trust_inherited_metadata", + } + # The sibling of SaveDatasetUnit: that one saves where the load said, this + # one is told where to save. + assert set(units["SaveDatasetToPathUnit"]["schema"]["properties"]) == {"path"} def test_component_fields_tell_the_front_which_components_to_offer(units): @@ -122,10 +152,17 @@ def test_component_fields_tell_the_front_which_components_to_offer(units): converter = units["ApplyConverterUnit"]["schema"]["properties"]["converter"] explorer = units["RunExplorationUnit"]["schema"]["properties"]["explorer"] + uploaded = units["LoadUploadedDatasetUnit"]["schema"]["properties"]["dataloader"] + datafile = units["LoadDatafileDatasetUnit"]["schema"]["properties"]["dataloader"] + assert model["parent"] == "BaseModel" assert optimizer["parent"] == "BaseOptimizer" assert converter["parent"] == "BaseConverter" assert explorer["parent"] == "BaseExplorer" + # Both loading units offer the same readers; what differs is how each one + # finds the bytes to hand them. + assert uploaded["parent"] == "BaseDataLoader" + assert datafile["parent"] == "BaseDataLoader" # Global and local explainers are separate registries with separate base # classes, and a component field carries a single parent hint. Hence two diff --git a/tests/back/units/test_dataset_ingest_units.py b/tests/back/units/test_dataset_ingest_units.py new file mode 100644 index 000000000..36796a5db --- /dev/null +++ b/tests/back/units/test_dataset_ingest_units.py @@ -0,0 +1,491 @@ +"""Contract tests for the units DatasetJob is built from. + +The context is built by hand rather than through the 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 pandas as pd +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.units.apply_dataset_schema_unit import ApplyDatasetSchemaUnit +from DashAI.back.units.compute_dataset_metadata_unit import ( + EXTENDED_METADATA_KEYS, + ComputeDatasetMetadataUnit, +) +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.infer_dataset_types_unit import InferDatasetTypesUnit +from DashAI.back.units.load_datafile_dataset_unit import LoadDatafileDatasetUnit +from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit + +SCHEMA = { + "n": {"type": "Float", "dtype": "float64"}, + "label": {"type": "Categorical", "dtype": "string"}, +} + + +@pytest.fixture(name="dataset") +def fixture_dataset(): + frame = pd.DataFrame({"n": [1.0, 2.0, 3.0], "label": ["a", "b", "a"]}) + return to_dashai_dataset(frame) + + +@pytest.fixture(name="ctx") +def fixture_ctx(dataset): + ctx = ExecutionContext() + ctx.put("dataset", dataset) + return ctx + + +# --------------------------------------------------------------------------- # +# ComputeDatasetMetadataUnit +# --------------------------------------------------------------------------- # + + +def test_metadata_unit_computes_the_extended_keys_by_default(ctx): + ComputeDatasetMetadataUnit(compute_metadata=True)(ctx) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key in splits + + +def test_metadata_unit_strips_extended_keys_it_did_not_ask_for(ctx): + """A dataset can arrive carrying metadata from wherever it was copied from. + + Asking for base only has to mean base only, whatever the file happened to + bring — otherwise the result depends on the dataset's history. + """ + dataset = ctx.require("dataset") + dataset.splits["correlations"] = {"stale": "value"} + dataset.splits["general_info"] = {"stale": "value"} + + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key not in splits + + +def test_metadata_unit_reuses_trusted_extended_metadata(ctx): + """Trusting what is there is the whole point of the flag: the marker value + below survives only because nothing was recomputed.""" + dataset = ctx.require("dataset") + for key in EXTENDED_METADATA_KEYS: + dataset.splits[key] = {"marker": key} + + ComputeDatasetMetadataUnit(compute_metadata=True, trust_inherited_metadata=True)( + ctx + ) + + splits = ctx.require("dataset").splits + assert splits["correlations"] == {"marker": "correlations"} + + +def test_metadata_unit_computes_when_there_is_nothing_to_trust(ctx): + """Trusting metadata that is not there would silently store none at all.""" + ComputeDatasetMetadataUnit(compute_metadata=True, trust_inherited_metadata=True)( + ctx + ) + + splits = ctx.require("dataset").splits + assert splits["total_rows"] == 3 + for key in EXTENDED_METADATA_KEYS: + assert key in splits + + +def test_metadata_unit_keeps_the_same_dataset_object(ctx): + """It fills the dataset in place, so downstream units that already hold a + reference see the metadata too. Asserted with ``is``: an equal copy would + pass ``==`` and still break that.""" + before = ctx.require("dataset") + + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + assert ctx.require("dataset") is before + + +def test_metadata_unit_refuses_to_run_without_a_dataset(): + with pytest.raises(UnitContractError): + ComputeDatasetMetadataUnit(compute_metadata=True)(ExecutionContext()) + + +# --------------------------------------------------------------------------- # +# InferDatasetTypesUnit +# --------------------------------------------------------------------------- # + + +def test_infer_types_unit_publishes_a_type_per_column(ctx): + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + inferred = ctx.require("inferred_types") + assert set(inferred) == {"n", "label"} + + +def test_infer_types_unit_prefers_the_types_the_dataset_already_carries(ctx): + """A reader that got types from the source knows better than inference over + the values, so those have to win.""" + typed = ApplyDatasetSchemaUnit(column_renames=None) + ctx.put_ref("inferred_types", SCHEMA) + typed(ctx) + + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + inferred = ctx.require("inferred_types") + assert inferred["label"]["type"] == "Categorical" + + +def test_infer_types_unit_publishes_something_json_serializable(ctx): + """It is a ref, not a live object: it has to survive ``put_ref``'s check.""" + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + + # Round-trips through the serializable half without raising. + assert "inferred_types" in ctx.to_dict() + + +# --------------------------------------------------------------------------- # +# ApplyDatasetSchemaUnit +# --------------------------------------------------------------------------- # + + +def test_apply_schema_unit_casts_the_declared_types(ctx): + ctx.put_ref("inferred_types", SCHEMA) + + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + types = ctx.require("dataset").types + assert type(types["label"]).__name__ == "Categorical" + + +def test_apply_schema_unit_carries_the_types_through_a_rename(ctx): + """The declared type has to follow the column, not the old name.""" + ctx.put_ref("inferred_types", SCHEMA) + + ApplyDatasetSchemaUnit(column_renames={"label": "variety"})(ctx) + + dataset = ctx.require("dataset") + assert dataset.column_names == ["n", "variety"] + assert type(dataset.types["variety"]).__name__ == "Categorical" + + +def test_apply_schema_unit_rejects_renames_that_collide(ctx): + ctx.put_ref("inferred_types", SCHEMA) + + with pytest.raises(JobError) as excinfo: + ApplyDatasetSchemaUnit(column_renames={"label": "n"})(ctx) + + assert "contain duplicates: ['n']" in str(excinfo.value) + + +def test_apply_schema_unit_rejects_a_declaration_for_columns_that_are_gone(ctx): + """The underlying transform passes unknown columns through untouched, so a + stale declaration would silently leave columns with the wrong types. This is + the check that turns that into an error. + """ + ctx.put_ref("inferred_types", dict(SCHEMA, removed_column={"type": "Float"})) + + with pytest.raises(JobError) as excinfo: + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + assert "removed_column" in str(excinfo.value) + + +def test_apply_schema_unit_validate_runs_before_any_work(ctx): + """``validate`` is a precondition check, so it must not need the transform to + have run — and ``__call__`` runs it on its own.""" + ctx.put_ref("inferred_types", {"not_a_column": {"type": "Float"}}) + + with pytest.raises(JobError): + ApplyDatasetSchemaUnit(column_renames=None).validate(ctx) + + +def test_apply_schema_unit_refuses_to_run_without_a_declaration(ctx): + """Declared in REQUIRES, so a missing value is a wiring mistake and has to + read as one instead of as "no types to apply".""" + with pytest.raises(UnitContractError): + ApplyDatasetSchemaUnit(column_renames=None)(ctx) + + +def test_infer_then_apply_compose_over_the_same_key(ctx): + """The pair is meant to chain: one publishes what the other consumes.""" + InferDatasetTypesUnit(method="DashAIPtype")(ctx) + ApplyDatasetSchemaUnit(column_renames={"n": "number"})(ctx) + + assert ctx.require("dataset").column_names == ["number", "label"] + + +# --------------------------------------------------------------------------- # +# SaveDatasetToPathUnit +# --------------------------------------------------------------------------- # + + +def test_save_to_path_unit_writes_where_it_is_told(ctx, tmp_path): + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + destination = tmp_path / "somewhere" / "dataset" + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + SaveDatasetToPathUnit(path=str(destination))(ctx) + + assert load_dataset(str(destination)).column_names == ["n", "label"] + + +def test_save_to_path_unit_ignores_dataset_path_entirely(ctx, tmp_path): + """Its whole reason to exist: a context left over from a load must not be + able to redirect the save back onto the source.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + source = tmp_path / "source" + destination = tmp_path / "destination" + ctx.put_ref("dataset_path", str(source)) + ComputeDatasetMetadataUnit(compute_metadata=False)(ctx) + + SaveDatasetToPathUnit(path=str(destination))(ctx) + + assert not source.exists() + assert load_dataset(str(destination)).column_names == ["n", "label"] + + +def test_save_to_path_unit_keeps_the_underlying_error_text(ctx, tmp_path): + """A save failure is an infrastructure failure, and only the outermost + message reaches the user — so the original text has to stay in it.""" + destination = tmp_path / "dataset" + + from DashAI.back.dataloaders.classes import dashai_dataset + + real_save = dashai_dataset.save_dataset + + def fail(dataset, path): + raise OSError("no space left on device") + + dashai_dataset.save_dataset = fail + try: + with pytest.raises(JobError) as excinfo: + SaveDatasetToPathUnit(path=str(destination))(ctx) + finally: + dashai_dataset.save_dataset = real_save + + message = str(excinfo.value) + assert str(destination) in message + assert "no space left on device" in message + + +# --------------------------------------------------------------------------- # +# the loading units +# --------------------------------------------------------------------------- # + + +@pytest.fixture(name="csv_file") +def fixture_csv_file(tmp_path): + path = tmp_path / "data.csv" + path.write_text("n,label\n1.0,a\n2.0,b\n", encoding="utf-8") + return path + + +@pytest.fixture(name="registry") +def fixture_registry(): + """A registry holding only the CSV reader, injected without an app.""" + from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader + from DashAI.back.dependencies.registry import ComponentRegistry + + registry = ComponentRegistry(initial_components=[CSVDataLoader]) + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def test_uploaded_unit_reads_the_file_and_publishes_the_dataset( + registry, csv_file, tmp_path +): + ctx = ExecutionContext() + + LoadUploadedDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_uploaded_unit_publishes_no_path_and_no_id(registry, csv_file, tmp_path): + """An uploaded file is not a stored dataset yet: there is no id to correlate + against and nowhere it belongs on disk. Publishing either would hand a later + unit a value that describes something else.""" + ctx = ExecutionContext() + + LoadUploadedDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert not ctx.has("dataset_path") + assert not ctx.has("dataset_id") + + +def test_uploaded_unit_reports_an_unknown_reader_with_the_registry_wording( + registry, csv_file, tmp_path +): + ctx = ExecutionContext() + + with pytest.raises(KeyError) as excinfo: + LoadUploadedDatasetUnit( + dataloader={"component": "NoSuchLoader", "params": {}}, + source=str(csv_file), + temp_path=str(tmp_path), + n_sample=None, + )(ctx) + + assert "does not exists in the registry" in str(excinfo.value) + + +class _DatafileRow: + def __init__(self, local_path, status): + self.local_path = local_path + self.status = status + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """A class, not a lambda: kink calls any registered lambda with the container + to resolve it, so a lambda would be invoked as a service factory instead of + handed to the unit as one.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +@pytest.fixture(name="datafile_rows") +def fixture_datafile_rows(): + rows = {} + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def test_datafile_unit_reads_the_selected_file( + registry, datafile_rows, csv_file, tmp_path +): + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + ctx = ExecutionContext() + + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + datafile_id=7, + selected_file="data.csv", + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_datafile_unit_skips_dotted_paths_when_choosing( + registry, datafile_rows, csv_file, tmp_path +): + """Download tools leave metadata directories behind that sort before the real + data, so without the filter the wrong file would win.""" + from DashAI.back.core.enums.status import DatafileStatus + + hidden = tmp_path / ".cache" + hidden.mkdir() + (hidden / "aaa.csv").write_text("junk\n", encoding="utf-8") + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + ctx = ExecutionContext() + + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {"separator": ","}}, + datafile_id=7, + selected_file=None, + )(ctx) + + assert ctx.require("dataset").column_names == ["n", "label"] + + +def test_datafile_unit_rejects_a_download_that_is_not_finished( + registry, datafile_rows, tmp_path +): + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.DOWNLOADING) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=7, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Datafile 7 is not ready." + + +def test_datafile_unit_treats_a_missing_row_the_same_way(registry, datafile_rows): + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=99, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Datafile 99 is not ready." + + +def test_datafile_unit_reports_an_empty_download(registry, datafile_rows, tmp_path): + from DashAI.back.core.enums.status import DatafileStatus + + empty = tmp_path / "empty" + empty.mkdir() + datafile_rows[7] = _DatafileRow(str(empty), DatafileStatus.READY) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "CSVDataLoader", "params": {}}, + datafile_id=7, + selected_file=None, + )(ExecutionContext()) + + assert str(excinfo.value) == "Hub download directory is empty." + + +def test_datafile_unit_has_its_own_wording_for_an_unknown_reader( + registry, datafile_rows, csv_file, tmp_path +): + """Deliberately different from the uploaded unit's, which surfaces the + registry's own ``KeyError``. Both texts reach users today.""" + from DashAI.back.core.enums.status import DatafileStatus + + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "NoSuchLoader", "params": {}}, + datafile_id=7, + selected_file="data.csv", + )(ExecutionContext()) + + assert str(excinfo.value) == "DataLoader 'NoSuchLoader' not found in registry." From 1b1ff8ee1f30b3b2026ed2f306ae168a3ccf8e38 Mon Sep 17 00:00:00 2001 From: Felipe Date: Thu, 13 Aug 2026 18:20:40 -0400 Subject: [PATCH 13/28] fix: Improve error messages for unknown explainers and data loaders --- DashAI/back/units/explanation_artifacts.py | 4 +-- .../back/units/load_datafile_dataset_unit.py | 10 ++++-- tests/back/api/test_explainer_job.py | 10 +++--- tests/back/units/test_dataset_ingest_units.py | 31 +++++++++++++++++++ 4 files changed, 46 insertions(+), 9 deletions(-) diff --git a/DashAI/back/units/explanation_artifacts.py b/DashAI/back/units/explanation_artifacts.py index 32608b742..e8ee8416a 100644 --- a/DashAI/back/units/explanation_artifacts.py +++ b/DashAI/back/units/explanation_artifacts.py @@ -55,8 +55,8 @@ def build_explainer(scope: str, selected: dict, trained_model: Any) -> Any: except Exception as e: log.exception(e) raise JobError( - f"""Unable to find the {scope} explainer with name - {explainer_name} in registry.""", + f"Unable to find the {scope} explainer with name " + f"{explainer_name} in registry.", ) from e try: diff --git a/DashAI/back/units/load_datafile_dataset_unit.py b/DashAI/back/units/load_datafile_dataset_unit.py index 2b97c4f28..b73631a9a 100644 --- a/DashAI/back/units/load_datafile_dataset_unit.py +++ b/DashAI/back/units/load_datafile_dataset_unit.py @@ -133,11 +133,15 @@ def execute(self, ctx: ExecutionContext) -> None: source = _resolve_source_file(work_dir, selected_file) + # Looked up among the readers specifically, not with the registry's + # global ``registry[name]``: that one walks every component type, so a + # metric or a model whose name happened to be passed here would be + # instantiated as if it could read a file instead of being rejected. dataloader_name = dataloader_config["component"] - registry = component_registry.registry.get("DataLoader", {}) - if dataloader_name not in registry: + readers = component_registry.registry.get("DataLoader", {}) + if dataloader_name not in readers: raise JobError(f"DataLoader '{dataloader_name}' not found in registry.") - dataloader = registry[dataloader_name]["class"]() + dataloader = readers[dataloader_name]["class"]() log.debug("Loading hub dataset from %s using %s", source, dataloader_name) ctx.put( diff --git a/tests/back/api/test_explainer_job.py b/tests/back/api/test_explainer_job.py index 2e55b422b..513751f51 100644 --- a/tests/back/api/test_explainer_job.py +++ b/tests/back/api/test_explainer_job.py @@ -503,8 +503,11 @@ def test_a_model_that_cannot_be_loaded_names_the_path(client, run_id): 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.""" + """One line, so it stays readable wherever it surfaces. + + It used to be a triple-quoted f-string, which put a newline and the source + file's indentation literally inside the text the user reads. + """ explainer_id = _create_global_explainer( client, run_id, explainer_name="NoSuchExplainer" ) @@ -513,8 +516,7 @@ def test_an_unknown_explainer_name_is_reported_with_the_multiline_message( 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." + "Unable to find the global explainer with name NoSuchExplainer in registry." ) diff --git a/tests/back/units/test_dataset_ingest_units.py b/tests/back/units/test_dataset_ingest_units.py index 36796a5db..bb2f86398 100644 --- a/tests/back/units/test_dataset_ingest_units.py +++ b/tests/back/units/test_dataset_ingest_units.py @@ -489,3 +489,34 @@ def test_datafile_unit_has_its_own_wording_for_an_unknown_reader( )(ExecutionContext()) assert str(excinfo.value) == "DataLoader 'NoSuchLoader' not found in registry." + + +def test_datafile_unit_rejects_a_component_that_is_not_a_reader( + datafile_rows, csv_file, tmp_path +): + """A registered component of some other kind is still not a reader. + + The registry's ``registry[name]`` indexer searches every type at once, so + looking the name up that way would find, say, a metric and call it as if it + could parse a file. The lookup is deliberately scoped to the readers. + """ + from DashAI.back.core.enums.status import DatafileStatus + from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader + from DashAI.back.dependencies.registry import ComponentRegistry + from DashAI.back.metrics.classification.accuracy import Accuracy + + di["component_registry"] = ComponentRegistry( + initial_components=[CSVDataLoader, Accuracy] + ) + datafile_rows[7] = _DatafileRow(str(tmp_path), DatafileStatus.READY) + try: + with pytest.raises(JobError) as excinfo: + LoadDatafileDatasetUnit( + dataloader={"component": "Accuracy", "params": {}}, + datafile_id=7, + selected_file="data.csv", + )(ExecutionContext()) + finally: + del di["component_registry"] + + assert str(excinfo.value) == "DataLoader 'Accuracy' not found in registry." From 3f648a66c5650abd7d4d47eb4cc10d760c3c0516 Mon Sep 17 00:00:00 2001 From: Felipe Date: Thu, 20 Aug 2026 15:00:07 -0400 Subject: [PATCH 14/28] test dag --- tests/back/pipeline_spike/__init__.py | 0 tests/back/pipeline_spike/dag_engine.py | 259 ++++++++++++ .../pipeline_spike/test_dag_engine_spike.py | 375 ++++++++++++++++++ 3 files changed, 634 insertions(+) create mode 100644 tests/back/pipeline_spike/__init__.py create mode 100644 tests/back/pipeline_spike/dag_engine.py create mode 100644 tests/back/pipeline_spike/test_dag_engine_spike.py diff --git a/tests/back/pipeline_spike/__init__.py b/tests/back/pipeline_spike/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/pipeline_spike/dag_engine.py b/tests/back/pipeline_spike/dag_engine.py new file mode 100644 index 000000000..b446ea085 --- /dev/null +++ b/tests/back/pipeline_spike/dag_engine.py @@ -0,0 +1,259 @@ +"""Minimal sequential DAG engine over the existing atomic units. + +A spike, not a deliverable. It exists to find out which parts of the unit +contract (``BaseUnit`` + ``ExecutionContext``) survive being driven by a graph +instead of by a job, before more jobs get atomized. It deliberately has no UI, +no database, no persistence, no API and no parallelism. + +The one design decision it puts to the test is the one proposed in section 9 of +ATOMIZING_JOBS.md: a graph engine does not need a shared context. It gives every +node an ``ExecutionContext`` of its own, preloaded with exactly the keys that +node declares in ``REQUIRES``, and the renaming lives on the edge rather than in +the unit. That is what lets two ``LoadDatasetUnit`` instances -- both of which +write the fixed key ``dataset`` -- coexist in one graph. + +Nothing here imports from ``DashAI/back/pipeline/``: that subsystem does not run +and is being redesigned. Nothing here modifies a unit either. Whatever the +engine cannot express with the contract as it stands is reported, not patched. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Sequence, Set, Tuple + +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + + +class GraphError(Exception): + """Raised when a graph is not executable as declared.""" + + +@dataclass(frozen=True) +class Node: + """A unit instance placed in a graph. + + Parameters + ---------- + id : str + Identifier of the node inside the graph. + unit : BaseUnit + The unit instance to run. Already configured. + seeds : Mapping[str, Any] + Constants injected straight into this node's context, for keys no + upstream unit publishes. ``run_id`` is the real case: four units + require it and none provides it, because in a job it arrived through + ``self.kwargs``. See the open question in the spike's report. + """ + + id: str + unit: BaseUnit + seeds: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class Edge: + """A single key travelling from one node's output to another's input. + + The pair of key names is what makes fixed ``REQUIRES``/``PROVIDES`` strings + behave as port names: ``src_key`` is drawn from the source's ``PROVIDES`` + and ``dst_key`` from the target's ``REQUIRES``, and the edge maps one onto + the other. + """ + + src: str + src_key: str + dst: str + dst_key: str + + +def connect(src: Node, dst: Node) -> Tuple[Edge, ...]: + """Wire every key the two nodes agree on: ``PROVIDES`` against ``REQUIRES``. + + A visual canvas cannot draw one edge per key: ``FitModelUnit`` alone + requires eight. This is the bundling rule a real engine would use -- one + drawn edge between two nodes stands for this whole set. + """ + shared = sorted(set(src.unit.PROVIDES) & set(dst.unit.REQUIRES)) + return tuple(Edge(src.id, key, dst.id, key) for key in shared) + + +@dataclass(frozen=True) +class Graph: + """A set of nodes and the edges between them.""" + + nodes: Sequence[Node] + edges: Sequence[Edge] + + def node(self, node_id: str) -> Node: + for node in self.nodes: + if node.id == node_id: + return node + raise GraphError(f"Unknown node id: {node_id}") + + +def validate(graph: Graph) -> List[str]: + """Check the graph statically and return a topological execution order. + + This is the static DAG validator ATOMIZING_JOBS.md lists as missing. It is + almost free: ``REQUIRES``/``PROVIDES`` already carry everything it needs, + and no unit has to run. + + Raises + ------ + GraphError + If any check fails, with every problem found listed at once. + """ + problems: List[str] = [] + ids = [node.id for node in graph.nodes] + if len(set(ids)) != len(ids): + problems.append("Duplicate node ids in the graph.") + + by_id = {node.id: node for node in graph.nodes} + inbound: Dict[str, Dict[str, List[Edge]]] = {node_id: {} for node_id in by_id} + + for edge in graph.edges: + if edge.src not in by_id or edge.dst not in by_id: + problems.append(f"Edge {edge} points at a node that is not in the graph.") + continue + if edge.src_key not in by_id[edge.src].unit.PROVIDES: + problems.append( + f"{edge.src} does not provide '{edge.src_key}' " + f"(provides: {list(by_id[edge.src].unit.PROVIDES)})." + ) + if edge.dst_key not in by_id[edge.dst].unit.REQUIRES: + problems.append( + f"{edge.dst} does not require '{edge.dst_key}' " + f"(requires: {list(by_id[edge.dst].unit.REQUIRES)})." + ) + inbound[edge.dst].setdefault(edge.dst_key, []).append(edge) + + for node in graph.nodes: + for key in node.seeds: + if key not in node.unit.REQUIRES: + problems.append( + f"{node.id} is seeded with '{key}', which it never uses." + ) + for key in node.unit.REQUIRES: + sources = len(inbound[node.id].get(key, [])) + (key in node.seeds) + if sources == 0: + problems.append(f"{node.id} requires '{key}' and nothing supplies it.") + elif sources > 1: + problems.append( + f"{node.id} gets '{key}' from {sources} sources; a port takes one." + ) + + order = _topological_order(graph, by_id, problems) + if problems: + raise GraphError("\n".join(problems)) + return order + + +def _topological_order( + graph: Graph, by_id: Mapping[str, Node], problems: List[str] +) -> List[str]: + """Kahn's algorithm. Leftover nodes mean a cycle.""" + pending = dict.fromkeys(by_id, 0) + successors: Dict[str, List[str]] = {node_id: [] for node_id in by_id} + for edge in graph.edges: + if edge.src in by_id and edge.dst in by_id: + pending[edge.dst] += 1 + successors[edge.src].append(edge.dst) + + ready = [node_id for node_id, count in pending.items() if count == 0] + order: List[str] = [] + while ready: + node_id = ready.pop(0) + order.append(node_id) + for successor in successors[node_id]: + pending[successor] -= 1 + if pending[successor] == 0: + ready.append(successor) + + if len(order) != len(by_id): + problems.append( + "The graph has a cycle: " + ", ".join(sorted(set(by_id) - set(order))) + ) + return order + + +#: How many times the engine had to deep-copy the whole reference half of a +#: context just to ask which half a single key lives in. See the report. +REF_PROBES = [0] + + +def _transport(src_ctx: ExecutionContext, key: str) -> Tuple[str, Any]: + """Read a value out of a context, keeping the half it came from. + + The engine has to move a value from one context to another, and the two + halves have incompatible rules: ``put_ref`` validates with ``json.dumps`` + and deep-copies, ``put`` stores the live object by reference. Handing a + ``DashAIDataset`` to ``put_ref`` raises; handing a dict meant as a reference + to ``put`` silently drops the copy-on-write guarantee. + + ``ExecutionContext`` has no public way to ask which half a key is in: + ``get``/``has`` merge them and ``refs``/``to_dict`` deep-copy the entire + reference half. So this probe is correct but expensive, and the cost is + counted rather than hidden. + """ + REF_PROBES[0] += 1 + if key in src_ctx.refs: + return "ref", src_ctx.get(key) + return "cache", src_ctx.require(key) + + +def run(graph: Graph) -> Dict[str, ExecutionContext]: + """Execute the graph sequentially and return each node's own context. + + Parameters + ---------- + graph : Graph + A graph that passes :func:`validate`. + + Returns + ------- + Dict[str, ExecutionContext] + The context of every node whose values were still needed at the end, + keyed by node id. Contexts nobody reads any more are dropped as the run + proceeds so intermediate datasets do not pile up. + """ + order = validate(graph) + by_id = {node.id: node for node in graph.nodes} + + consumers: Dict[str, int] = dict.fromkeys(by_id, 0) + for edge in graph.edges: + consumers[edge.src] += 1 + + contexts: Dict[str, ExecutionContext] = {} + for node_id in order: + node = by_id[node_id] + ctx = ExecutionContext() + + for key, value in node.seeds.items(): + ctx.put_ref(key, value) + + for edge in graph.edges: + if edge.dst != node_id: + continue + half, value = _transport(contexts[edge.src], edge.src_key) + if half == "ref": + ctx.put_ref(edge.dst_key, value) + else: + ctx.put(edge.dst_key, value) + + node.unit(ctx) + contexts[node_id] = ctx + + for edge in graph.edges: + if edge.dst != node_id: + continue + consumers[edge.src] -= 1 + if consumers[edge.src] == 0: + contexts.pop(edge.src).clear_cache() + + return contexts + + +def sinks(graph: Graph) -> Set[str]: + """Node ids nothing consumes. They are what a run exists to produce.""" + consumed = {edge.src for edge in graph.edges} + return {node.id for node in graph.nodes} - consumed diff --git a/tests/back/pipeline_spike/test_dag_engine_spike.py b/tests/back/pipeline_spike/test_dag_engine_spike.py new file mode 100644 index 000000000..c74079077 --- /dev/null +++ b/tests/back/pipeline_spike/test_dag_engine_spike.py @@ -0,0 +1,375 @@ +"""Runs real units as a DAG, to find out what the unit contract cannot express. + +The graph below is the train/test converter flow, which is the smallest thing a +single shared context genuinely cannot run: two ``LoadDatasetUnit`` instances +both write the fixed key ``dataset``, so in one context the second load erases +the first. Section 9 of ATOMIZING_JOBS.md proposes giving each node a context of +its own and putting the renaming on the edge; these tests check whether that +actually works with the units exactly as they are. + + load_train --dataset--+-------------------> tx_train --dataset--> save_train + | ^ ^ + +-> fit --converter--+---+ dataset_path | + | --+ + load_test --dataset------------------------|-> tx_test --dataset--> save_test + +Two fan-outs (``load_train`` feeds two nodes, ``fit`` feeds two nodes) and two +joins (each ``tx`` node takes a dataset from one branch and a fitted converter +from another). The number that proves it worked is 2.0: MinMaxScaler fitted on +train [0, 5, 10] learns min=0 max=10, so a test value of 20 scales to 2.0. A +refit on the test data would have produced 0.0 instead. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.types.value_types import Float +from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit +from tests.back.pipeline_spike.dag_engine import ( + Edge, + Graph, + GraphError, + Node, + connect, + run, + sinks, + validate, +) + +FULL_SCOPE = {"columns": [], "rows": []} +_MIN_MAX = {"component": "MinMaxScaler", "params": {}} + + +class _Row: + """Stand-in for a Dataset ORM row, as in tests/back/units.""" + + def __init__(self, file_path): + self.file_path = file_path + self.dataset_id = None + + +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: + """A class, not a lambda: kink resolves a registered lambda by calling it.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +def _write(root, **columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +def _values(path, column="a"): + return list(load_dataset(str(path)).to_pandas()[column]) + + +@pytest.fixture(name="stored_datasets") +def fixture_stored_datasets(tmp_path): + """Dataset 7 is the training data, dataset 8 the test data.""" + train = _write(tmp_path / "train", a=[0.0, 5.0, 10.0]) + test = _write(tmp_path / "test", a=[20.0]) + + di["session_factory"] = _FakeSessionFactory( + {"Dataset": {7: _Row(str(train)), 8: _Row(str(test))}} + ) + di["component_registry"] = {"MinMaxScaler": {"class": MinMaxScaler}} + yield train, test + del di["session_factory"] + del di["component_registry"] + + +def _graph(tmp_path, converter="MinMaxScaler"): + """The diamond, wired with the bundling rule wherever it applies.""" + load_train = Node("load_train", LoadDatasetUnit(dataset_id=7)) + load_test = Node("load_test", LoadDatasetUnit(dataset_id=8)) + fit = Node( + "fit", + FitConverterUnit( + converter={"component": converter, "params": {}}, + scope=FULL_SCOPE, + target=None, + ), + ) + tx_train = Node("tx_train", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + tx_test = Node("tx_test", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + # Saves back over the training data, so it needs the ref `dataset_path` from + # the load as well as the live dataset from the transform: a join whose two + # inputs live in different halves of the context. + save_train = Node("save_train", SaveDatasetUnit()) + save_test = Node( + "save_test", SaveDatasetToPathUnit(path=str(tmp_path / "out" / "dataset")) + ) + + nodes = [load_train, load_test, fit, tx_train, tx_test, save_train, save_test] + edges = [ + *connect(load_train, fit), + Edge("load_train", "dataset", "tx_train", "dataset"), + Edge("load_train", "dataset_path", "save_train", "dataset_path"), + *connect(fit, tx_train), + *connect(fit, tx_test), + Edge("load_test", "dataset", "tx_test", "dataset"), + Edge("tx_train", "dataset", "save_train", "dataset"), + *connect(tx_test, save_test), + ] + return Graph(nodes, edges) + + +def test_the_diamond_runs_and_the_fitted_state_survives_the_branch( + stored_datasets, tmp_path +): + """The headline case: one fit, two branches, no refit. + + 2.0 on the test branch is only reachable if the converter fitted in the + ``fit`` node arrived at ``tx_test`` still carrying the range it learned from + the training data, having crossed two context boundaries on the way. + """ + train, _test = stored_datasets + + run(_graph(tmp_path)) + + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + assert _values(tmp_path / "out" / "dataset") == [2.0] + + +def test_two_loads_of_the_same_unit_do_not_fight_over_the_key( + stored_datasets, tmp_path +): + """Section 9's proposal, checked directly. + + Both loads publish the fixed key ``dataset``. In one shared context the + second would overwrite the first and the test branch would silently scale + the training data. + """ + graph = _graph(tmp_path) + contexts = run(graph) + + assert sinks(graph) == {"save_train", "save_test"} + # Every context still alive at the end belongs to a sink; the loads and the + # fit were dropped as soon as nothing needed them. + assert set(contexts) == {"save_train", "save_test"} + + +def test_the_engine_never_needed_a_shared_key_to_be_renamed_by_a_unit(tmp_path): + """The edge does the renaming; the units keep their fixed key names.""" + graph = _graph(tmp_path) + + assert set(validate(graph)[:2]) == {"load_train", "load_test"} + assert all(edge.src_key == edge.dst_key for edge in graph.edges) + + +def test_the_bundling_rule_picks_the_keys_two_nodes_agree_on(): + """One drawn edge between two nodes stands for a set of keys.""" + load = Node("load", LoadDatasetUnit(dataset_id=1)) + save = Node("save", SaveDatasetUnit()) + + assert {edge.src_key for edge in connect(load, save)} == {"dataset", "dataset_path"} + + +def test_a_missing_input_is_caught_before_anything_runs(): + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + load = Node("load", LoadDatasetUnit(dataset_id=1)) + + with pytest.raises(GraphError, match="tx requires 'fitted_converter'"): + validate(Graph([load, tx], list(connect(load, tx)))) + + +def test_two_edges_into_one_port_are_rejected(): + """The failure the old engine had: a merge where the last writer wins.""" + load_a = Node("a", LoadDatasetUnit(dataset_id=1)) + load_b = Node("b", LoadDatasetUnit(dataset_id=2)) + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + graph = Graph([load_a, load_b, fit], [*connect(load_a, fit), *connect(load_b, fit)]) + + with pytest.raises(GraphError, match="fit gets 'dataset' from 2 sources"): + validate(graph) + + +def test_a_cycle_is_reported(): + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + graph = Graph([fit, tx], [*connect(fit, tx), *connect(tx, fit)]) + + with pytest.raises(GraphError, match="cycle"): + validate(graph) + + +def test_run_id_has_to_be_seeded_because_no_unit_publishes_it(): + """``run_id`` is an orphan input: four units require it, none provides it. + + In a job it arrived through ``self.kwargs``. A graph has to inject it from + outside, and the validator cannot tell that injection apart from a wire the + user forgot to draw. + """ + save = Node("save", SaveModelUnit()) + + with pytest.raises(GraphError, match="save requires 'run_id'"): + validate(Graph([save], [])) + + seeded = Node("save", SaveModelUnit(), seeds={"run_id": 1}) + with pytest.raises(GraphError, match="save requires 'model'"): + validate(Graph([seeded], [])) + + +def test_a_seed_for_a_key_the_unit_never_uses_is_rejected(): + save = Node("save", SaveModelUnit(), seeds={"run_id": 1, "nonsense": 2}) + + with pytest.raises(GraphError, match="seeded with 'nonsense'"): + validate(Graph([save], [])) + + +def test_an_optional_output_cannot_be_wired_at_all(): + """``PROVIDES`` is the whole vocabulary an edge can name. + + ``FitModelUnit`` produces ``best_parameters`` only on the hyperparameter + search branch, so it cannot declare it (``__call__`` checks ``PROVIDES`` + unconditionally). ``ModelJob`` copes by reading it with ``ctx.has(...)``. + A graph cannot: a key outside ``PROVIDES`` is not addressable by an edge, + so the limitation goes from "not verifiable" to "not connectable". + """ + fit = Node( + "fit", + FitConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + tx = Node("tx", TransformDatasetUnit(scope=FULL_SCOPE, target=None)) + graph = Graph([fit, tx], [Edge("fit", "best_parameters", "tx", "dataset")]) + + with pytest.raises(GraphError, match="fit does not provide 'best_parameters'"): + validate(graph) + + +def test_a_key_a_middle_node_does_not_republish_needs_an_edge_around_it(): + """A converter in the middle does not carry the ids the load published. + + ``ApplyConverterUnit`` provides only ``dataset`` and ``fitted_converter``, + and it is right not to republish ``dataset_id`` -- section 4.3 forbids + anything derived from the object being transformed from crossing the + boundary. But ``PrepareAndSplitUnit`` downstream needs both, so ``dataset`` + comes through the converter and ``dataset_id`` has to jump over it. + + The engine accepts the bypass. What it costs is on the canvas: the user has + to draw an arrow from the loader past the converters to the split. + """ + load = Node("load", LoadDatasetUnit(dataset_id=1)) + apply_ = Node( + "apply", + ApplyConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), + ) + split = Node("split", PrepareAndSplitUnit(splits={})) + + bundled = Graph( + [load, apply_, split], [*connect(load, apply_), *connect(apply_, split)] + ) + with pytest.raises(GraphError, match="split requires 'dataset_id'"): + validate(bundled) + + bypassed = Graph( + [load, apply_, split], + [ + *connect(load, apply_), + *connect(apply_, split), + Edge("load", "dataset_id", "split", "dataset_id"), + ], + ) + assert validate(bypassed) == ["load", "apply", "split"] + + +class _CountingConverter: + """Records how many times it was asked to transform, and who built it.""" + + instances = [] + CHANGES_ROW_COUNT = False + + def __init__(self, **params): + _CountingConverter.instances.append(self) + self.transform_calls = 0 + + def fit(self, x, y=None): + return self + + def transform(self, x, y=None): + self.transform_calls += 1 + return x + + +def test_a_fan_out_hands_both_branches_the_same_live_object(stored_datasets, tmp_path): + """Identity survives the fan-out, which it has to (ATOMIZING_JOBS 5.7). + + Some converters cache against the object they were given, so a copy would + silently recompute. It also means two consumers of one fitted converter + share mutable state -- harmless while the engine is sequential, and a + reason to keep it that way. + """ + _CountingConverter.instances.clear() + di["component_registry"]["Counting"] = {"class": _CountingConverter} + + graph = _graph(tmp_path, converter="Counting") + run(graph) + + assert len(_CountingConverter.instances) == 1 + assert _CountingConverter.instances[0].transform_calls == 2 + + +def test_moving_a_value_across_an_edge_keeps_the_half_it_came_from( + stored_datasets, tmp_path +): + """Refs travel as refs, live objects as live objects. + + ``dataset_path`` is a ref and ``dataset`` is a cached object, and the two + have incompatible rules: ``put_ref`` validates with ``json.dumps``, so a + dataset sent that way raises, and ``put`` would drop the copy-on-write + guarantee a ref depends on. The engine has to ask which half a key is in -- + and the only public way to ask deep-copies the entire reference half, once + per edge. + """ + from tests.back.pipeline_spike import dag_engine + + graph = _graph(tmp_path) + dag_engine.REF_PROBES[0] = 0 + contexts = run(graph) + + assert dag_engine.REF_PROBES[0] == len(graph.edges) + save_train = contexts["save_train"] + assert "dataset_path" in save_train.refs + assert "dataset" not in save_train.refs + assert save_train.has("dataset") From e02817715782146bc8d55afa4b5f11370aa1f9b0 Mon Sep 17 00:00:00 2001 From: Felipe Date: Fri, 21 Aug 2026 23:54:57 -0400 Subject: [PATCH 15/28] Add comprehensive tests for DAG tracking, unit evaluation, and model saving - Introduced tests for tracking DAG execution in `test_tracking.py`, ensuring proper recording of node runs and artifacts. - Updated `test_dag_engine_spike.py` to reflect changes in run_id handling, ensuring it is no longer an orphan input. - Enhanced `test_build_model_unit.py` with tests for handling run_id and model evaluation metrics. - Added tests in `test_evaluate_model_to_artifact_unit.py` to verify model evaluation and metric publication. - Implemented validation checks in `test_evaluate_model_unit.py` to ensure run_id presence during evaluation. - Created `test_save_model_unit.py` to validate model saving behavior and artifact prefix handling. - Updated `test_fit_model_unit.py` to ensure proper handling of model state and artifact naming conventions. - Enhanced unit contract tests in `test_unit_contracts.py` to ensure configuration keys are properly declared and classified. --- .../e7b4d1a9c206_add_pipeline_run_tracking.py | 122 ++++ DashAI/back/core/enums/status.py | 21 + DashAI/back/dag/__init__.py | 28 + DashAI/back/dag/engine.py | 184 ++++++ DashAI/back/dag/expand.py | 217 +++++++ DashAI/back/dag/graph.py | 180 ++++++ DashAI/back/dag/tracking.py | 155 +++++ DashAI/back/dag/validate.py | 199 +++++++ DashAI/back/dependencies/database/models.py | 163 ++++++ DashAI/back/initial_components.py | 4 + DashAI/back/job/model_job.py | 13 +- DashAI/back/job/pipeline_job.py | 262 +++++---- DashAI/back/models/base_model.py | 110 +++- DashAI/back/optimizers/base_optimizer.py | 27 +- DashAI/back/units/base_unit.py | 17 + DashAI/back/units/build_manual_input_unit.py | 26 +- DashAI/back/units/build_model_unit.py | 18 +- .../units/compute_dataset_metadata_unit.py | 30 +- DashAI/back/units/context.py | 30 + .../units/evaluate_model_to_artifact_unit.py | 90 +++ DashAI/back/units/evaluate_model_unit.py | 35 +- DashAI/back/units/fit_model_unit.py | 40 +- .../units/generate_local_explanation_unit.py | 27 +- .../back/units/load_training_dataset_unit.py | 37 +- .../back/units/load_uploaded_dataset_unit.py | 23 +- DashAI/back/units/save_model_unit.py | 38 +- tests/back/api/test_model_job_as_a_graph.py | 535 ++++++++++++++++++ tests/back/api/test_units_api.py | 76 ++- tests/back/dag/__init__.py | 0 tests/back/dag/test_engine.py | 474 ++++++++++++++++ tests/back/dag/test_expand.py | 285 ++++++++++ tests/back/dag/test_pipeline_job.py | 339 +++++++++++ tests/back/dag/test_tracking.py | 266 +++++++++ .../pipeline_spike/test_dag_engine_spike.py | 32 +- tests/back/units/test_build_model_unit.py | 40 +- tests/back/units/test_context.py | 39 ++ .../test_evaluate_model_to_artifact_unit.py | 156 +++++ tests/back/units/test_evaluate_model_unit.py | 33 +- tests/back/units/test_fit_model_unit.py | 95 +++- tests/back/units/test_save_model_unit.py | 91 +++ tests/back/units/test_unit_contracts.py | 257 +++++++++ 41 files changed, 4478 insertions(+), 336 deletions(-) create mode 100644 DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py create mode 100644 DashAI/back/dag/__init__.py create mode 100644 DashAI/back/dag/engine.py create mode 100644 DashAI/back/dag/expand.py create mode 100644 DashAI/back/dag/graph.py create mode 100644 DashAI/back/dag/tracking.py create mode 100644 DashAI/back/dag/validate.py create mode 100644 DashAI/back/units/evaluate_model_to_artifact_unit.py create mode 100644 tests/back/api/test_model_job_as_a_graph.py create mode 100644 tests/back/dag/__init__.py create mode 100644 tests/back/dag/test_engine.py create mode 100644 tests/back/dag/test_expand.py create mode 100644 tests/back/dag/test_pipeline_job.py create mode 100644 tests/back/dag/test_tracking.py create mode 100644 tests/back/units/test_evaluate_model_to_artifact_unit.py create mode 100644 tests/back/units/test_save_model_unit.py diff --git a/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py b/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py new file mode 100644 index 000000000..fdbeb079e --- /dev/null +++ b/DashAI/alembic/versions/e7b4d1a9c206_add_pipeline_run_tracking.py @@ -0,0 +1,122 @@ +"""Add pipeline run tracking tables. + +Separates the definition of a graph from its executions. ``Pipeline`` keeps the +definition, which changes as the user edits it; ``pipeline_run`` freezes the +steps and edges of one execution, ``pipeline_node_run`` tracks each node of it, +and ``pipeline_node_artifact`` holds what a node emitted, keyed by a key from +the unit's PROVIDES rather than by a column per node type. + +The three JSON result columns on ``pipeline`` (exploration, train, prediction) +are deliberately left in place: the pipelines endpoints and the front's results +view still read them, so removing them belongs with rewriting those. + +Revision ID: e7b4d1a9c206 +Revises: d5b3c8f2a041 +Create Date: 2026-08-20 00:00:00.000000 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +revision: str = "e7b4d1a9c206" +down_revision: Union[str, None] = "d5b3c8f2a041" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "pipeline_run", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("pipeline_id", sa.Integer(), nullable=False), + sa.Column("steps", sa.JSON(), nullable=True), + sa.Column("edges", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("last_modified", sa.DateTime(), nullable=False), + sa.Column("delivery_time", sa.DateTime(), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "status", + sa.Enum( + "NOT_STARTED", + "DELIVERED", + "STARTED", + "FINISHED", + "ERROR", + name="pipelinerunstatus", + ), + nullable=False, + ), + sa.Column("error_message", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["pipeline_id"], + ["pipeline.id"], + name=op.f("fk_pipeline_run_pipeline_id_pipeline"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_run")), + ) + + op.create_table( + "pipeline_node_run", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("pipeline_run_id", sa.Integer(), nullable=False), + sa.Column("node_id", sa.String(), nullable=False), + sa.Column("block_id", sa.String(), nullable=False), + sa.Column("node_type", sa.String(), nullable=False), + sa.Column("config", sa.JSON(), nullable=True), + sa.Column("input", sa.JSON(), nullable=True), + sa.Column("output", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.Column("last_modified", sa.DateTime(), nullable=False), + sa.Column("delivery_time", sa.DateTime(), nullable=True), + sa.Column("start_time", sa.DateTime(), nullable=True), + sa.Column("end_time", sa.DateTime(), nullable=True), + sa.Column( + "status", + sa.Enum( + "NOT_STARTED", + "DELIVERED", + "STARTED", + "FINISHED", + "ERROR", + "CANCELLED", + name="noderunstatus", + ), + nullable=False, + ), + sa.Column("error_message", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["pipeline_run_id"], + ["pipeline_run.id"], + name=op.f("fk_pipeline_node_run_pipeline_run_id_pipeline_run"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_node_run")), + ) + + op.create_table( + "pipeline_node_artifact", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("node_run_id", sa.Integer(), nullable=False), + sa.Column("key", sa.String(), nullable=False), + sa.Column("value", sa.JSON(), nullable=True), + sa.Column("created", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["node_run_id"], + ["pipeline_node_run.id"], + name=op.f("fk_pipeline_node_artifact_node_run_id_pipeline_node_run"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_pipeline_node_artifact")), + ) + + +def downgrade() -> None: + op.drop_table("pipeline_node_artifact") + op.drop_table("pipeline_node_run") + op.drop_table("pipeline_run") diff --git a/DashAI/back/core/enums/status.py b/DashAI/back/core/enums/status.py index a21973878..12b90f08f 100644 --- a/DashAI/back/core/enums/status.py +++ b/DashAI/back/core/enums/status.py @@ -59,3 +59,24 @@ class DatafileStatus(Enum): DOWNLOADING = "downloading" READY = "ready" ERROR = "error" + + +class PipelineRunStatus(Enum): + NOT_STARTED = 0 + DELIVERED = 1 + STARTED = 2 + FINISHED = 3 + ERROR = 4 + + +class NodeRunStatus(Enum): + NOT_STARTED = 0 + DELIVERED = 1 + STARTED = 2 + FINISHED = 3 + ERROR = 4 + # A node that never ran because an earlier one failed. Distinct from + # NOT_STARTED, which is a node still waiting its turn: without the + # distinction a run that died halfway is indistinguishable from one still + # in flight. + CANCELLED = 5 diff --git a/DashAI/back/dag/__init__.py b/DashAI/back/dag/__init__.py new file mode 100644 index 000000000..f144f1d12 --- /dev/null +++ b/DashAI/back/dag/__init__.py @@ -0,0 +1,28 @@ +"""A sequential DAG engine over the atomic units. + +Nothing here imports from ``DashAI/back/pipeline/``. That subsystem does not +run — its nodes implement two of ``BaseJob``'s four abstract methods, so they +cannot even be instantiated — and it is replaced rather than repaired. + +The engine is deliberately sequential. Its predecessor was concurrent, and the +concurrency was not buying what it cost: the most expensive node in a pipeline +was already serialised behind an exclusive lock, every context and database +write was wrapped in another, and SQLite answered the remaining concurrent +writes with "database is locked" often enough to need retries with exponential +backoff. Units open a database session each, which is safe in sequence and was +the source of that contention in parallel. +""" + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, connect, sinks +from DashAI.back.dag.validate import resolve_unit_class, validate + +__all__ = [ + "Edge", + "Graph", + "GraphError", + "Node", + "connect", + "resolve_unit_class", + "sinks", + "validate", +] diff --git a/DashAI/back/dag/engine.py b/DashAI/back/dag/engine.py new file mode 100644 index 000000000..2d1a3a221 --- /dev/null +++ b/DashAI/back/dag/engine.py @@ -0,0 +1,184 @@ +"""Sequential execution of a validated graph. + +The engine gives every node an ``ExecutionContext`` of its own, preloaded with +exactly the keys that node declares in ``REQUIRES``, and the renaming lives on +the edge. There is no shared context: merging the outputs of two predecessors +into one dictionary loses a key whenever both publish the same one, which is +how its predecessor lost values silently and why it grew a parallel list of raw +branch dictionaries to work around itself. + +Nothing here touches the database. Tracking goes through a sink, so the engine +can run with no persistence at all — which is how it is tested. +""" + +import logging +from typing import Any, Dict, List, Mapping, Optional, Protocol + +from DashAI.back.dag.graph import Graph, GraphError +from DashAI.back.dag.validate import instantiate, resolve_unit_class, validate +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class TrackingSink(Protocol): + """Where a run reports what it is doing. + + Every method is a notification, not a decision: a sink may persist, log, or + do nothing, and the engine's behaviour does not depend on which. + """ + + def run_started(self, order: List[str]) -> None: + """The graph passed validation and is about to run, in this order.""" + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + """``payload`` is the serializable half of the node's input context.""" + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + """``artifacts`` are the node's serializable outputs, by PROVIDES key.""" + + def node_failed(self, node_id: str, message: str) -> None: ... + + def nodes_cancelled(self, node_ids: List[str]) -> None: + """Nodes that will never run because an earlier one failed.""" + + def run_finished(self) -> None: ... + + def run_failed(self, message: str) -> None: ... + + +class NullSink: + """A sink that records nothing. The engine's default.""" + + def run_started(self, order: List[str]) -> None: + pass + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + pass + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + pass + + def node_failed(self, node_id: str, message: str) -> None: + pass + + def nodes_cancelled(self, node_ids: List[str]) -> None: + pass + + def run_finished(self) -> None: + pass + + def run_failed(self, message: str) -> None: + pass + + +def artifacts_of(unit_class: type, ctx: ExecutionContext) -> Dict[str, Any]: + """The outputs of a node that can be persisted, by ``PROVIDES`` key. + + Only the reference half: the cache holds live datasets, models and tasks, + which are not serializable and are always derivable again. A unit whose + real output is on disk publishes the path as a reference, so the path is + what gets recorded — which is the whole of what a caller needs later. + + ``origin`` is what makes this cheap. Asking ``key in ctx.refs`` answers the + same question but deep-copies every reference the context holds, once per + key. + """ + return { + key: ctx.get(key) for key in unit_class.PROVIDES if ctx.origin(key) == "ref" + } + + +def run( + graph: Graph, sink: Optional[TrackingSink] = None +) -> Dict[str, ExecutionContext]: + """Execute a graph sequentially, returning the contexts still in use. + + Parameters + ---------- + graph : Graph + The graph to run. It is validated first, so a graph that cannot work + fails before any unit does. + sink : Optional[TrackingSink] + Where progress is reported. Defaults to recording nothing. + + Returns + ------- + Dict[str, ExecutionContext] + The context of every node whose values were still needed at the end, + keyed by node id. Contexts nobody reads any more are dropped as the run + proceeds, so intermediate datasets do not pile up for the length of the + run. + + Raises + ------ + GraphError + If the graph does not validate. + """ + sink = sink if sink is not None else NullSink() + + order = validate(graph) + by_id = graph.by_id() + classes = {node.id: resolve_unit_class(node.unit) for node in graph.nodes} + + # One instance per node, built before anything runs: a configuration that + # cannot build its unit is a mistake in the graph, not a failure of a node + # halfway through it. + units = {node.id: instantiate(node) for node in graph.nodes} + + consumers: Dict[str, int] = dict.fromkeys(by_id, 0) + for edge in graph.edges: + consumers[edge.src] += 1 + + sink.run_started(order) + + contexts: Dict[str, ExecutionContext] = {} + for position, node_id in enumerate(order): + ctx = ExecutionContext() + + for edge in graph.edges: + if edge.dst != node_id: + continue + source = contexts[edge.src] + # The half a value lives in decides how it may be moved: put_ref + # validates with json.dumps and deep-copies, put stores the live + # object by reference. A dataset handed to put_ref raises, and a + # reference handed to put loses its copy-on-read guarantee. + if source.origin(edge.src_key) == "ref": + ctx.put_ref(edge.dst_key, source.get(edge.src_key)) + else: + ctx.put(edge.dst_key, source.require(edge.src_key)) + + sink.node_started(node_id, ctx.to_dict()) + try: + units[node_id](ctx) + except Exception as e: + log.exception(e) + sink.node_failed(node_id, str(e)) + # Everything after this point in the order never runs. Reporting + # them keeps a run that died halfway distinguishable from one still + # in flight. + sink.nodes_cancelled(list(order[position + 1 :])) + sink.run_failed(str(e)) + raise + + contexts[node_id] = ctx + sink.node_finished(node_id, artifacts_of(classes[node_id], ctx), ctx.to_dict()) + + # Release what nothing downstream will read again. + for edge in graph.edges: + if edge.dst != node_id: + continue + consumers[edge.src] -= 1 + if consumers[edge.src] == 0 and edge.src in contexts: + contexts.pop(edge.src).clear_cache() + + sink.run_finished() + return contexts + + +__all__ = ["GraphError", "NullSink", "TrackingSink", "artifacts_of", "run"] diff --git a/DashAI/back/dag/expand.py b/DashAI/back/dag/expand.py new file mode 100644 index 000000000..f5f300183 --- /dev/null +++ b/DashAI/back/dag/expand.py @@ -0,0 +1,217 @@ +"""Turning what a canvas holds into the graph the engine runs. + +A canvas cannot show a training as six nodes and fifteen edges, so a block on +it stands for a sequence of units. The engine works on units, so a run is over +the expanded graph, and it is the expanded graph a run freezes: replaying one +never depends on re-running the expansion that produced it. + +The first version has one unit per block, and expansion is close to a rename. +The shape is here from the start because the persistence has to be right before +there is anything to migrate. +""" + +import re +from typing import Any, Dict, List, Mapping, Sequence + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node +from DashAI.back.dag.validate import resolve_unit_class + +#: Anything outside what can name a directory under RUNS_PATH, plus the escape +#: character itself. ``SaveModelUnit.validate`` refuses the rest, so a prefix +#: built from a node id a user chose has to be brought into that alphabet here. +#: +#: ``_`` is in here on purpose. Replacing unsafe characters with a fixed ``_`` +#: is not injective -- ``save.a`` and ``save_a`` would collapse to the same +#: name, and the second saving node would overwrite the first one's model +#: directory in silence. Escaping to ``_`` instead, with ``_`` itself +#: escaped, is reversible, so two different node ids cannot produce one prefix. +_NEEDS_ESCAPE = re.compile(r"[^A-Za-z0-9-]") + +#: The runtime params this module knows how to answer. A unit declares which +#: ones it takes in its own ``RUNTIME_PARAMS``; these are the two that are about +#: the pipeline run rather than about a job, so a pipeline is what supplies them. +#: Any other runtime param belongs to whoever else runs that unit. +ENGINE_SUPPLIED = ("artifact_prefix", "run_id") + + +def artifact_prefix(pipeline_run_id: int, node_id: str) -> str: + """Name this node's artifacts so nothing else can collide with them. + + The runs directory is shared with every real ``Run``, and a pipeline run id + is a different sequence from a run id: both start at 1, so a pipeline that + used its own id as the name would write over the model of the run with that + id. Not a risk -- a certainty. The ``pipeline-`` prefix is what keeps the + two apart, the run id keeps two executions of the same pipeline apart, and + the node id keeps two saving nodes in one graph apart. + + The node id is escaped rather than cleaned, so the mapping is injective and + two different ids can never produce one prefix. + """ + escaped = _NEEDS_ESCAPE.sub( + lambda match: f"_{ord(match.group()):02x}", str(node_id) + ) + return f"pipeline-{pipeline_run_id}-{escaped}" + + +def expand( + steps: Sequence[Mapping[str, Any]], + edges: Sequence[Mapping[str, Any]], + pipeline_run_id: int, +) -> Graph: + """Expand a canvas into a unit-level graph. + + Parameters + ---------- + steps : Sequence[Mapping[str, Any]] + The blocks, each ``{id, units: [{id, unit, config}], ...}``. + edges : Sequence[Mapping[str, Any]] + The wires between blocks, each ``{source, target}``. + pipeline_run_id : int + The run being expanded, which is what names its artifacts. + + Returns + ------- + Graph + Nodes are units, and every node carries the block it came from. + + Raises + ------ + GraphError + If a block is malformed, or if it holds more than one unit — see the + note below. + + Notes + ----- + A block holding several units is refused rather than guessed at. Wiring + across a boundary where either side is a sequence has more than one + defensible answer, and the rule that bundles keys does not cover every wire + a real graph needs: ``ApplyConverterUnit`` deliberately does not republish + ``dataset_id``, so a node needing both the converted dataset and the id + takes the id on an edge that skips the converter. Choosing a rule for that + without a case to check it against would be inventing semantics. The + persistence already carries N units per block, so adding it later costs no + migration. + """ + blocks = _read_blocks(steps) + + nodes: List[Node] = [] + for block_id, units in blocks.items(): + for unit in units: + nodes.append( + Node( + id=unit["id"], + unit=unit["unit"], + config=_with_engine_config( + unit["unit"], + unit.get("config") or {}, + pipeline_run_id, + unit["id"], + ), + block_id=block_id, + ) + ) + + wires: List[Edge] = [] + for edge in edges: + source, target = edge.get("source"), edge.get("target") + if source not in blocks or target not in blocks: + raise GraphError( + f"An edge connects '{source}' to '{target}', and one of them is " + "not a block in this pipeline." + ) + wires.extend(_between(blocks[source][-1], blocks[target][0])) + + return Graph(nodes, wires) + + +def _read_blocks( + steps: Sequence[Mapping[str, Any]], +) -> Dict[str, List[Mapping[str, Any]]]: + """Validate the shape of the blocks and index them by id.""" + blocks: Dict[str, List[Mapping[str, Any]]] = {} + + for step in steps: + block_id = step.get("id") + if not block_id: + raise GraphError("A block has no id.") + if block_id in blocks: + raise GraphError(f"There is more than one block with the id '{block_id}'.") + + units = step.get("units") + if not units: + raise GraphError( + f"Block '{block_id}' declares no units. Blocks saved by the " + "previous pipeline subsystem name a node type and carry a " + "single config instead, and cannot be run by this engine." + ) + if len(units) > 1: + raise GraphError( + f"Block '{block_id}' holds {len(units)} units. Only one unit " + "per block is supported so far." + ) + for unit in units: + if not unit.get("id") or not unit.get("unit"): + raise GraphError( + f"A unit in block '{block_id}' is missing its id or its unit name." + ) + + blocks[block_id] = list(units) + + return blocks + + +def _with_engine_config( + unit_name: str, + config: Mapping[str, Any], + pipeline_run_id: int, + node_id: str, +) -> Dict[str, Any]: + """Set the configuration the engine owns rather than the user. + + Whatever the stored graph holds for these fields is **discarded**, not + merged. Neither is a value a user could get right: + + * ``artifact_prefix`` is built from the id of the run, and there is no run + when a node is configured -- so any value already there was chosen + without the one thing that decides it. Overriding is also what makes two + nodes sharing a prefix impossible rather than merely unlikely, so the + collision needs no validator to catch it. + * ``run_id`` is null because a pipeline has no ``Run`` row. Forcing it is + what makes the sandbox mechanical instead of conventional: a stored graph + cannot point a training node at a real run and start writing ``Metric`` + rows into it. + + Both are declared in the unit's ``RUNTIME_PARAMS`` rather than in its + schema, so a form never offers them in the first place. This is the other + half of the same statement, on the side a hand-edited row could still reach. + """ + resolved = dict(config) + takes = set(resolve_unit_class(unit_name).RUNTIME_PARAMS) + + if "artifact_prefix" in takes: + resolved["artifact_prefix"] = artifact_prefix(pipeline_run_id, node_id) + + if "run_id" in takes: + resolved["run_id"] = None + + return resolved + + +def _between(source: Mapping[str, Any], target: Mapping[str, Any]) -> List[Edge]: + """Wire the keys the two units agree on: ``PROVIDES`` against ``REQUIRES``. + + One drawn edge stands for the whole set, because a canvas cannot draw one + per key. + """ + provides = set(resolve_unit_class(source["unit"]).PROVIDES) + requires = set(resolve_unit_class(target["unit"]).REQUIRES) + shared = sorted(provides & requires) + + if not shared: + raise GraphError( + f"'{source['id']}' ({source['unit']}) has nothing " + f"'{target['id']}' ({target['unit']}) needs, so the edge between " + "them would carry nothing." + ) + + return [Edge(source["id"], key, target["id"], key) for key in shared] diff --git a/DashAI/back/dag/graph.py b/DashAI/back/dag/graph.py new file mode 100644 index 000000000..7ee2cd75a --- /dev/null +++ b/DashAI/back/dag/graph.py @@ -0,0 +1,180 @@ +"""The shape of a graph: nodes, edges, and the rule that bundles them.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, Mapping, Optional, Sequence, Set, Tuple + + +class GraphError(Exception): + """Raised when a graph is not executable as declared.""" + + +@dataclass(frozen=True) +class Node: + """One unit placed in a graph. + + A node carries the *name* of its unit rather than an instance. Unit + instances hold state of their own — ``FitModelUnit`` memoizes its resolved + optimizer, ``FitConverterUnit`` its converter class — so two nodes sharing + one instance would have the second silently run with the first one's. The + engine builds one instance per node from ``(unit, config)``, which is also + exactly what a persisted graph can hold. + + Parameters + ---------- + id : str + Identifier of this node inside the graph. + unit : str + Class name of the unit, as the component registry knows it. + config : Mapping[str, Any] + The unit's own configuration, validated by its ``SCHEMA``. Constants a + caller chooses live here — not in the context, where nothing upstream + could ever supply them. + block_id : Optional[str] + The visual block this node belongs to, for an editor that groups + several units into one shape on a canvas. Left out, it becomes ``id``: + a node always belongs to some block, so nothing downstream — the + tracking column least of all — has a null to handle, and a node has one + representation rather than two. + """ + + id: str + unit: str + config: Mapping[str, Any] = field(default_factory=dict) + block_id: Optional[str] = None + + def __post_init__(self) -> None: + if self.block_id is None: + object.__setattr__(self, "block_id", self.id) + + +@dataclass(frozen=True) +class Edge: + """One key travelling from a node's output to another's input. + + The pair of names is what makes the fixed strings in ``REQUIRES`` and + ``PROVIDES`` behave as port names: ``src_key`` is drawn from the source's + ``PROVIDES`` and ``dst_key`` from the target's ``REQUIRES``, and the edge + maps one onto the other. That is what lets two ``LoadDatasetUnit`` nodes, + which both write the fixed key ``dataset``, coexist in one graph — the + renaming lives here rather than in the unit. + """ + + src: str + src_key: str + dst: str + dst_key: str + + +@dataclass(frozen=True) +class Graph: + """A set of nodes and the edges between them.""" + + nodes: Sequence[Node] + edges: Sequence[Edge] + + def node(self, node_id: str) -> Node: + for node in self.nodes: + if node.id == node_id: + return node + raise GraphError(f"Unknown node id: {node_id}") + + def by_id(self) -> Dict[str, Node]: + return {node.id: node for node in self.nodes} + + +def connect(src: Node, dst: Node) -> Tuple[Edge, ...]: + """Wire every key two nodes agree on: ``PROVIDES`` against ``REQUIRES``. + + A canvas cannot draw one edge per key — ``FitModelUnit`` alone requires + seven — so one drawn edge between two nodes stands for this whole set. + + The rule does not cover every wire a real graph needs, and that is by + design rather than an omission: ``ApplyConverterUnit`` does not republish + ``dataset_id``, because nothing derived from the object being transformed + may cross the context boundary. A downstream node that needs both the + converted dataset and the id therefore takes the id on an edge that skips + the converter entirely. + """ + from DashAI.back.dag.validate import resolve_unit_class + + provides = set(resolve_unit_class(src.unit).PROVIDES) + requires = set(resolve_unit_class(dst.unit).REQUIRES) + return tuple(Edge(src.id, key, dst.id, key) for key in sorted(provides & requires)) + + +def sinks(graph: Graph) -> Set[str]: + """Node ids nothing consumes. They are what a run exists to produce.""" + consumed = {edge.src for edge in graph.edges} + return {node.id for node in graph.nodes} - consumed + + +def dump(graph: Graph) -> Tuple[list, list]: + """Serialize a graph to the plain data a run freezes. + + Unit level, not block level: a run records the graph it actually executed, + so replaying it never depends on re-running the expansion that produced it. + """ + steps = [ + { + "id": node.id, + "block_id": node.block_id, + "unit": node.unit, + "config": dict(node.config), + } + for node in graph.nodes + ] + edges = [ + { + "src": edge.src, + "src_key": edge.src_key, + "dst": edge.dst, + "dst_key": edge.dst_key, + } + for edge in graph.edges + ] + return steps, edges + + +def load( + steps: Sequence[Mapping[str, Any]], edges: Sequence[Mapping[str, Any]] +) -> Graph: + """Rebuild a graph from what :func:`dump` produced. + + Raises + ------ + GraphError + If a step or an edge is missing a field. Rows written by the previous + subsystem land here: their steps carry a node ``type`` and a single + ``config``, with no unit to resolve, so they are refused with a message + that says so rather than failing somewhere deeper. + """ + try: + nodes = [ + Node( + id=step["id"], + unit=step["unit"], + config=step.get("config") or {}, + block_id=step.get("block_id"), + ) + for step in steps + ] + except KeyError as e: + raise GraphError( + f"A step is missing the field {e}. Steps saved by the previous " + "pipeline subsystem name a node type rather than a unit and cannot " + "be run by this engine." + ) from e + + try: + wires = [ + Edge(edge["src"], edge["src_key"], edge["dst"], edge["dst_key"]) + for edge in edges + ] + except KeyError as e: + raise GraphError( + f"An edge is missing the field {e}. Edges saved by the previous " + "pipeline subsystem connect nodes without naming the keys they " + "carry, so they cannot be run by this engine." + ) from e + + return Graph(nodes, wires) diff --git a/DashAI/back/dag/tracking.py b/DashAI/back/dag/tracking.py new file mode 100644 index 000000000..929079b52 --- /dev/null +++ b/DashAI/back/dag/tracking.py @@ -0,0 +1,155 @@ +"""Persistence of a run: the sink the engine reports to. + +Kept apart from the engine on purpose. The engine decides what happens and in +what order; this decides what is written down about it, and the engine runs +identically with a sink that writes nothing. + +The run row is created by the caller, not here. Naming a run's artifacts needs +its id, and that has to happen before the graph is built rather than after, so +the caller creates the row, expands with its id, and hands the id over. This +writes progress into it and nothing else -- a graph that fails to validate +leaves the row exactly as the caller made it, and saying why is the caller's +job. +""" + +import logging +from typing import Any, Dict, List, Mapping + +from DashAI.back.dag.graph import Graph, dump + +log = logging.getLogger(__name__) + + +class DatabaseSink: + """Writes a run to ``PipelineRun``, ``NodeRun`` and ``NodeArtifact``. + + Each notification opens its own short session and commits, so the canvas + can colour a node the moment it changes rather than at the end of the run. + That is safe here and was not in the concurrent predecessor, which needed a + process-wide lock and retries with backoff to survive SQLite answering + "database is locked". + + Parameters + ---------- + pipeline_run_id : int + The run to write into. Already created by the caller. + graph : Graph + The expanded, unit-level graph. It is frozen into the run, so the run + stays readable after the pipeline it came from is edited. + """ + + def __init__(self, pipeline_run_id: int, graph: Graph) -> None: + self.pipeline_run_id = pipeline_run_id + self._graph = graph + self._nodes = graph.by_id() + #: node id -> NodeRun id, so later notifications need no lookup by name. + self._node_run_ids: Dict[str, int] = {} + + @property + def _session_factory(self): + from kink import di + + return di["session_factory"] + + def run_started(self, order: List[str]) -> None: + """Freeze the graph into the run, and create a row per node. + + Every node gets a row up front, in ``NOT_STARTED``. Creating them on + demand instead would leave a node that never ran with no row at all, + and "no row" cannot distinguish a node still waiting from one whose run + died before reaching it. + """ + from DashAI.back.dependencies.database.models import NodeRun, PipelineRun + + steps, edges = dump(self._graph) + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.steps = steps + pipeline_run.edges = edges + pipeline_run.set_status_as_started() + db.commit() + + node_runs = [ + NodeRun( + pipeline_run_id=self.pipeline_run_id, + node_id=node_id, + block_id=self._nodes[node_id].block_id, + node_type=self._nodes[node_id].unit, + config=dict(self._nodes[node_id].config), + ) + for node_id in order + ] + db.add_all(node_runs) + db.commit() + self._node_run_ids = {row.node_id: row.id for row in node_runs} + + def node_started(self, node_id: str, payload: Mapping[str, Any]) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + with self._session_factory() as db: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_started() + node_run.input = dict(payload) + db.commit() + + def node_finished( + self, node_id: str, artifacts: Mapping[str, Any], payload: Mapping[str, Any] + ) -> None: + """Record the node as finished, with one artifact row per output key. + + The key is a key from the unit's own ``PROVIDES``. That is what + replaced a column per node type: adding a kind of node used to mean + adding a column to the pipeline table and a branch to an ``if + node_type`` chain. + """ + from DashAI.back.dependencies.database.models import NodeArtifact, NodeRun + + node_run_id = self._node_run_ids[node_id] + with self._session_factory() as db: + node_run = db.get(NodeRun, node_run_id) + node_run.set_status_as_finished() + node_run.output = dict(payload) + db.add_all( + [ + NodeArtifact(node_run_id=node_run_id, key=key, value=value) + for key, value in artifacts.items() + ] + ) + db.commit() + + def node_failed(self, node_id: str, message: str) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + with self._session_factory() as db: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_error(message) + db.commit() + + def nodes_cancelled(self, node_ids: List[str]) -> None: + from DashAI.back.dependencies.database.models import NodeRun + + if not node_ids: + return + + with self._session_factory() as db: + for node_id in node_ids: + node_run = db.get(NodeRun, self._node_run_ids[node_id]) + node_run.set_status_as_cancelled() + db.commit() + + def run_finished(self) -> None: + from DashAI.back.dependencies.database.models import PipelineRun + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.set_status_as_finished() + db.commit() + + def run_failed(self, message: str) -> None: + from DashAI.back.dependencies.database.models import PipelineRun + + with self._session_factory() as db: + pipeline_run = db.get(PipelineRun, self.pipeline_run_id) + pipeline_run.set_status_as_error(message) + db.commit() diff --git a/DashAI/back/dag/validate.py b/DashAI/back/dag/validate.py new file mode 100644 index 000000000..435bf7ae7 --- /dev/null +++ b/DashAI/back/dag/validate.py @@ -0,0 +1,199 @@ +"""Static validation of a graph, and the registry lookup it rests on. + +Everything here runs before any unit does. ``REQUIRES`` and ``PROVIDES`` are +class attributes, so the whole check needs the unit *classes* and never an +instance: a graph that cannot work is rejected without a dataset being read or +a model being built. +""" + +from typing import Dict, List, Mapping + +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node + + +def resolve_unit_class(name: str) -> type: + """Look a unit class up in the component registry by name. + + Parameters + ---------- + name : str + Class name of the unit. + + Returns + ------- + type + The unit class. Not an instance: this is what static validation reads + ``REQUIRES`` and ``PROVIDES`` off, and building an instance would mean + running the unit's own configuration checks too early. + + Raises + ------ + GraphError + If the name is unknown, or names something that is not a unit. + """ + from kink import di + + from DashAI.back.units.base_unit import BaseUnit + + component_registry = di["component_registry"] + + try: + unit_class = component_registry[name]["class"] + except Exception as e: + raise GraphError(f"There is no unit named '{name}'.") from e + + if not (isinstance(unit_class, type) and issubclass(unit_class, BaseUnit)): + raise GraphError( + f"'{name}' is registered but it is not a unit, so it cannot be a " + "node in a pipeline." + ) + + return unit_class + + +def instantiate(node: Node): + """Build this node's own unit instance. + + One instance per node, never shared: a unit keeps state on itself — a + memoized optimizer, a resolved converter class — so two nodes sharing an + instance would have the second silently run with the first one's. + + Raises + ------ + GraphError + If the configuration does not build the unit. + """ + unit_class = resolve_unit_class(node.unit) + try: + return unit_class(**dict(node.config)) + except Exception as e: + raise GraphError( + f"Node '{node.id}' could not be configured as a {node.unit}: {e}" + ) from e + + +def validate(graph: Graph) -> List[str]: + """Check the graph and return a topological execution order. + + Every problem found is reported at once rather than one per run, because a + user fixing a graph on a canvas wants the whole list. + + Two kinds of missing input are caught here: a context key no edge feeds, + and a runtime param nobody supplied. + + There is no escape hatch for an unfed input. Every key in ``REQUIRES`` has + a unit somewhere that publishes it — the audit in + ``tests/back/units/test_unit_contracts.py`` enforces that — so "nothing + supplies this key" means a missing edge and nothing else. That was not true + while ``run_id`` was a context key no unit published: an injected constant + and a wire the user forgot to draw were indistinguishable. + + Raises + ------ + GraphError + If any check fails, listing every problem found. + """ + problems: List[str] = [] + + ids = [node.id for node in graph.nodes] + duplicates = sorted({node_id for node_id in ids if ids.count(node_id) > 1}) + if duplicates: + problems.append(f"Duplicate node ids in the graph: {duplicates}.") + + by_id = graph.by_id() + + # Resolve every unit up front: a name that is not a unit makes every + # contract check below meaningless, so it is reported on its own. + classes: Dict[str, type] = {} + for node in graph.nodes: + try: + classes[node.id] = resolve_unit_class(node.unit) + except GraphError as e: + problems.append(f"Node '{node.id}': {e}") + + if problems and not classes: + raise GraphError("\n".join(problems)) + + inbound: Dict[str, Dict[str, List[Edge]]] = {node_id: {} for node_id in by_id} + + for edge in graph.edges: + if edge.src not in by_id or edge.dst not in by_id: + problems.append(f"Edge {edge} points at a node that is not in the graph.") + continue + if edge.src in classes and edge.src_key not in classes[edge.src].PROVIDES: + problems.append( + f"'{edge.src}' does not provide '{edge.src_key}' " + f"(provides: {list(classes[edge.src].PROVIDES)})." + ) + if edge.dst in classes and edge.dst_key not in classes[edge.dst].REQUIRES: + problems.append( + f"'{edge.dst}' does not require '{edge.dst_key}' " + f"(requires: {list(classes[edge.dst].REQUIRES)})." + ) + inbound[edge.dst].setdefault(edge.dst_key, []).append(edge) + + for node in graph.nodes: + if node.id not in classes: + continue + # A runtime param nobody supplied is a KeyError halfway through the + # run otherwise -- after earlier nodes already wrote to disk. It also + # tells a user something they need to know: a unit whose runtime params + # only a particular job knows how to fill is not usable as a node yet. + missing = [ + param + for param in classes[node.id].RUNTIME_PARAMS + if param not in node.config + ] + if missing: + problems.append( + f"'{node.id}' ({node.unit}) needs {missing} supplied by " + "whatever runs it, and nothing in this pipeline supplies them." + ) + for key in classes[node.id].REQUIRES: + sources = len(inbound[node.id].get(key, [])) + if sources == 0: + problems.append( + f"'{node.id}' requires '{key}' and no edge supplies it." + ) + elif sources > 1: + problems.append( + f"'{node.id}' gets '{key}' from {sources} edges; a port " + "takes one. Merging two values into one input is the " + "ambiguity this rejects rather than resolving silently." + ) + + order = _topological_order(graph, by_id, problems) + + if problems: + raise GraphError("\n".join(problems)) + + return order + + +def _topological_order( + graph: Graph, by_id: Mapping[str, Node], problems: List[str] +) -> List[str]: + """Kahn's algorithm. Nodes left over are a cycle, found before running.""" + pending = dict.fromkeys(by_id, 0) + successors: Dict[str, List[str]] = {node_id: [] for node_id in by_id} + for edge in graph.edges: + if edge.src in by_id and edge.dst in by_id: + pending[edge.dst] += 1 + successors[edge.src].append(edge.dst) + + ready = [node_id for node_id, count in pending.items() if count == 0] + order: List[str] = [] + while ready: + node_id = ready.pop(0) + order.append(node_id) + for successor in successors[node_id]: + pending[successor] -= 1 + if pending[successor] == 0: + ready.append(successor) + + if len(order) != len(by_id): + problems.append( + "The graph has a cycle: " + ", ".join(sorted(set(by_id) - set(order))) + ) + + return order diff --git a/DashAI/back/dependencies/database/models.py b/DashAI/back/dependencies/database/models.py index 07741791c..84ad4dca2 100644 --- a/DashAI/back/dependencies/database/models.py +++ b/DashAI/back/dependencies/database/models.py @@ -28,6 +28,8 @@ DatasetStatus, ExplainerStatus, ExplorerStatus, + NodeRunStatus, + PipelineRunStatus, PluginStatus, PredictionStatus, RunStatus, @@ -563,10 +565,171 @@ class Pipeline(Base): name: Mapped[str] = mapped_column(String, nullable=False) steps: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) edges: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + # Dead columns, kept for compatibility. Results live in NodeArtifact now, + # keyed by a PROVIDES key rather than by node type. They are still read by + # the pipelines endpoints (dataexploration/results, filter_models) and by + # the results view in the front, so dropping them is a change to those. exploration: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) train: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) prediction: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + pipeline_runs: Mapped[List["PipelineRun"]] = relationship( + "PipelineRun", cascade="all, delete-orphan", back_populates="pipeline" + ) + + +class PipelineRun(Base): + __tablename__ = "pipeline_run" + """ + Table to store one execution of a pipeline. + + The definition of a graph lives in ``Pipeline`` and changes as the user + edits it; a run freezes the ``steps`` and ``edges`` it actually executed, + so a past execution stays readable after the pipeline it came from was + rewritten. + """ + id: Mapped[int] = mapped_column(primary_key=True) + pipeline_id: Mapped[int] = mapped_column( + ForeignKey("pipeline.id", ondelete="CASCADE"), nullable=False + ) + steps: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + edges: Mapped[List[Dict[str, Any]]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, default=datetime.now, onupdate=datetime.now + ) + delivery_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + start_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + end_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + status: Mapped[Enum] = mapped_column( + Enum(PipelineRunStatus), + nullable=False, + default=PipelineRunStatus.NOT_STARTED, + ) + error_message: Mapped[str] = mapped_column(String, nullable=True) + + pipeline: Mapped["Pipeline"] = relationship( + "Pipeline", back_populates="pipeline_runs" + ) + node_runs: Mapped[List["NodeRun"]] = relationship( + "NodeRun", cascade="all, delete-orphan", back_populates="pipeline_run" + ) + + def set_status_as_delivered(self) -> None: + """Update the status of the pipeline run to delivered.""" + self.status = PipelineRunStatus.DELIVERED + self.delivery_time = datetime.now() + + def set_status_as_started(self) -> None: + """Update the status of the pipeline run to started.""" + self.status = PipelineRunStatus.STARTED + self.start_time = datetime.now() + + def set_status_as_finished(self) -> None: + """Update the status of the pipeline run to finished.""" + self.status = PipelineRunStatus.FINISHED + self.end_time = datetime.now() + + def set_status_as_error(self, error_message: Optional[str] = None) -> None: + """Update the status of the pipeline run to error.""" + self.status = PipelineRunStatus.ERROR + self.error_message = error_message + self.end_time = datetime.now() + + +class NodeRun(Base): + __tablename__ = "pipeline_node_run" + """ + Table to store the execution of one node of a pipeline run. + + A node is one unit. ``block_id`` is the visual block the node belongs to: + an editor groups several units into one block on the canvas, so a block + maps to N node runs and its status is an aggregate of theirs. In the first + version every block holds exactly one unit and ``block_id`` equals + ``node_id``, but the column is here from the start so growing to N never + needs a migration. + """ + id: Mapped[int] = mapped_column(primary_key=True) + pipeline_run_id: Mapped[int] = mapped_column( + ForeignKey("pipeline_run.id", ondelete="CASCADE"), nullable=False + ) + node_id: Mapped[str] = mapped_column(String, nullable=False) + block_id: Mapped[str] = mapped_column(String, nullable=False) + #: Class name of the unit this node runs, as the component registry knows it. + node_type: Mapped[str] = mapped_column(String, nullable=False) + config: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + input: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + output: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + last_modified: Mapped[DateTime] = mapped_column( + DateTime, default=datetime.now, onupdate=datetime.now + ) + delivery_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + start_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + end_time: Mapped[DateTime] = mapped_column(DateTime, nullable=True) + status: Mapped[Enum] = mapped_column( + Enum(NodeRunStatus), + nullable=False, + default=NodeRunStatus.NOT_STARTED, + ) + error_message: Mapped[str] = mapped_column(String, nullable=True) + + pipeline_run: Mapped["PipelineRun"] = relationship( + "PipelineRun", back_populates="node_runs" + ) + artifacts: Mapped[List["NodeArtifact"]] = relationship( + "NodeArtifact", cascade="all, delete-orphan", back_populates="node_run" + ) + + def set_status_as_delivered(self) -> None: + """Update the status of the node run to delivered.""" + self.status = NodeRunStatus.DELIVERED + self.delivery_time = datetime.now() + + def set_status_as_started(self) -> None: + """Update the status of the node run to started.""" + self.status = NodeRunStatus.STARTED + self.start_time = datetime.now() + + def set_status_as_finished(self) -> None: + """Update the status of the node run to finished.""" + self.status = NodeRunStatus.FINISHED + self.end_time = datetime.now() + + def set_status_as_error(self, error_message: Optional[str] = None) -> None: + """Update the status of the node run to error.""" + self.status = NodeRunStatus.ERROR + self.error_message = error_message + self.end_time = datetime.now() + + def set_status_as_cancelled(self) -> None: + """Update the status of the node run to cancelled. + + A node that never ran because an earlier one failed, as opposed to one + still waiting its turn. + """ + self.status = NodeRunStatus.CANCELLED + self.end_time = datetime.now() + + +class NodeArtifact(Base): + __tablename__ = "pipeline_node_artifact" + """ + Table to store one output of a node run. + + ``key`` is a key from the unit's ``PROVIDES``, so what a node emits is + named by its own declared contract rather than by a column per node type. + """ + id: Mapped[int] = mapped_column(primary_key=True) + node_run_id: Mapped[int] = mapped_column( + ForeignKey("pipeline_node_run.id", ondelete="CASCADE"), nullable=False + ) + key: Mapped[str] = mapped_column(String, nullable=False) + value: Mapped[Dict[str, Any]] = mapped_column(JSON, nullable=True) + created: Mapped[DateTime] = mapped_column(DateTime, default=datetime.now) + + node_run: Mapped["NodeRun"] = relationship("NodeRun", back_populates="artifacts") + class Converter(Base): __tablename__ = "converter" diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index bd507979b..092948fcb 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -354,6 +354,9 @@ from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.compute_dataset_metadata_unit import ComputeDatasetMetadataUnit +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_converter_unit import FitConverterUnit from DashAI.back.units.fit_model_unit import FitModelUnit @@ -549,6 +552,7 @@ def get_initial_components(): BuildModelUnit, FitModelUnit, EvaluateModelUnit, + EvaluateModelToArtifactUnit, SaveModelUnit, ApplyConverterUnit, FitConverterUnit, diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 23e2c8bcb..24cbddf82 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -98,7 +98,9 @@ def run( # Get the necessary parameters run_id: int = self.kwargs["run_id"] - ctx = ExecutionContext(refs={"run_id": run_id}) + # The run id is not context: no unit publishes it, so it travels as + # configuration to each unit that needs it. + ctx = ExecutionContext() with session_factory() as db: run: Run = db.get(Run, run_id) @@ -132,6 +134,7 @@ def run( train_metrics=model_session.train_metrics, validation_metrics=model_session.validation_metrics, test_metrics=model_session.test_metrics, + run_id=run_id, )(ctx) # Resolving the optimizer before the status changes keeps an @@ -142,6 +145,10 @@ def run( "params": run.optimizer_parameters, }, goal_metric=run.goal_metric, + run_id=run_id, + # The run names its own artifacts, which is what keeps the + # plot filenames of two runs apart inside RUNS_PATH. + artifact_prefix=str(run_id), ) fit_model.validate(ctx) @@ -180,10 +187,10 @@ def run( ) from e self.report_progress(0.85, "Computing metrics") - EvaluateModelUnit()(ctx) + EvaluateModelUnit(run_id=run_id)(ctx) self.report_progress(0.95, "Saving model") - SaveModelUnit()(ctx) + SaveModelUnit(artifact_prefix=str(run_id))(ctx) try: run.run_path = ctx.require("model_path") diff --git a/DashAI/back/job/pipeline_job.py b/DashAI/back/job/pipeline_job.py index dacdfb81b..cd1dff625 100644 --- a/DashAI/back/job/pipeline_job.py +++ b/DashAI/back/job/pipeline_job.py @@ -1,114 +1,180 @@ +"""Job that runs a pipeline as a graph of units.""" + import logging -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING -from kink import di +from kink import inject -from DashAI.back.dependencies.database.models import Pipeline +from DashAI.back.dag.engine import run as run_graph +from DashAI.back.dag.expand import expand +from DashAI.back.dag.graph import GraphError +from DashAI.back.dag.tracking import DatabaseSink from DashAI.back.job.base_job import BaseJob, JobError if TYPE_CHECKING: - from sqlalchemy.orm import Session - - from DashAI.back.dependencies.registry import ComponentRegistry + from sqlalchemy.orm import sessionmaker log = logging.getLogger(__name__) class PipelineJob(BaseJob): + """Run a pipeline: expand its blocks into units and execute the graph. + + The job owns what a job owns -- the database session, the run row, its + status transitions -- and the engine owns execution. Between them sits the + expansion, which needs the run's id: naming a node's artifacts has to + happen before the graph is built, so the row is created here first and the + engine never creates anything. + + ``kwargs`` is a single id, as every job's is: the whole job is serialized + with dill to reach the worker process, which rebuilds its dependencies from + a fresh container, so nothing but plain data can travel in it. The graph is + read from the database inside ``run``. + """ + + @staticmethod + def _pipeline_id(kwargs) -> int: + """The id of the pipeline to run, under either name it arrives as. + + The front sends ``{"id": …}`` -- the wire contract predates this job -- + while every sibling job names its own subject (``run_id``, + ``converter_id``). Both are accepted so the existing caller keeps + working and a new one can use the name that matches the others; the + fallback goes when the endpoint is rewritten. + """ + pipeline_id = kwargs.get("pipeline_id", kwargs.get("id")) + if pipeline_id is None: + raise JobError("No pipeline id was given to run. Send it as 'pipeline_id'.") + return pipeline_id + def set_status_as_delivered(self) -> None: - pass + """Nothing to mark: the run row does not exist until the job starts. - async def run( - self, - component_registry: "ComponentRegistry" = lambda di: di["component_registry"], - ) -> None: - db: "Session" = self.kwargs["db"] - id: int = self.kwargs.get("id", None) - pipeline: Pipeline = db.get(Pipeline, id) - steps: List[Dict[str, Any]] = self.kwargs.get("steps", []) or pipeline.steps - - if not id: - raise JobError("No id provided to execute the pipeline.") - if not steps: - raise JobError("No steps provided to execute the pipeline.") - - if not steps: - raise JobError("Pipeline has no steps to execute") - - log.info(f"Starting pipeline execution for pipeline {id}...") - - context: Dict[str, Any] = {"pipeline_id": id} - - for idx, step in enumerate(steps): - node_id = step.get("id") - node_type = step.get("type") - node_config = step.get("config", {}) - - log.debug( - f"Pipeline {id}: Executing step {idx + 1}/{len(steps)} - " - f"{node_type} ({node_id})" - ) - - try: - node_class = component_registry(di)[node_type]["class"] - except KeyError as e: - error_msg = f"Component type {node_type} not found in registry" - raise JobError( - f"Error in node {node_id} ({node_type}): {error_msg}" - ) from e - - try: - node_instance = node_class(**node_config) - except Exception as e: - error_msg = f"Error in node {node_id} ({node_type}): {str(e)}" - log.exception(error_msg) - raise JobError(error_msg) from e - - try: - output = await node_instance.run(context=context) - self._update_context(context, pipeline, node_type, node_id, output) - log.debug(f"Node {node_id} executed successfully.") - - except Exception as e: - error_msg = f"Error in node {node_id} ({node_type}): {str(e)}" - log.exception(error_msg) - raise JobError(error_msg) from e - - log.info(f"Pipeline {id} execution completed successfully.") - db.add(pipeline) - db.commit() - self.set_status_as_delivered() - - def _update_context( + A pipeline has no per-pipeline status, and the run row is created in + ``run`` because expanding the graph needs its id. Recording the + delivered state would mean creating the row when the job is enqueued + instead, which belongs with the endpoint that enqueues it. + """ + log.debug("PipelineJob delivered; no row to mark yet.") + + def set_status_as_error(self) -> None: + """Nothing to mark. + + This is called on a job deleted while still queued, before ``run`` has + created anything. A run that has started and then failed is marked by + the sink, from inside ``run``. + """ + log.debug("PipelineJob errored before starting; no row to mark.") + + @inject + def get_job_name(self) -> str: + """Get a descriptive name for the job.""" + from kink import di + + from DashAI.back.dependencies.database.models import Pipeline + + try: + pipeline_id = self._pipeline_id(self.kwargs) + except JobError: + return "Pipeline" + + try: + with di["session_factory"]() as db: + pipeline = db.get(Pipeline, pipeline_id) + if pipeline and pipeline.name: + return f"Pipeline: {pipeline.name}" + except Exception: + pass + + return f"Pipeline ({pipeline_id})" + + @inject + def run( self, - context: Dict[str, Any], - pipeline: Pipeline, - node_type: str, - node_id: str, - output: Dict[str, Any], + session_factory: "sessionmaker" = lambda di: di["session_factory"], ) -> None: + import gc + + from DashAI.back.dependencies.database.models import Pipeline, PipelineRun + + pipeline_id: int = self._pipeline_id(self.kwargs) + + with session_factory() as db: + pipeline: Pipeline = db.get(Pipeline, pipeline_id) + if not pipeline: + raise JobError(f"Pipeline {pipeline_id} does not exist in DB.") + + steps = pipeline.steps or [] + edges = pipeline.edges or [] + if not steps: + raise JobError(f"Pipeline {pipeline_id} has no steps to run.") + + # Created before the graph is built, because naming a node's + # artifacts needs the run's id: two executions of one pipeline that + # shared a name would have the second write over the first's model. + pipeline_run = PipelineRun(pipeline_id=pipeline_id) + db.add(pipeline_run) + db.commit() + pipeline_run_id = pipeline_run.id + + self.report_progress(0.05, "Preparing the graph") + + try: + graph = expand(steps, edges, pipeline_run_id) + except GraphError as e: + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + + sink = DatabaseSink(pipeline_run_id, graph) + + try: + contexts = run_graph(graph, sink) + except GraphError as e: + # Validation happens before the engine reports anything, so the run + # is still untouched and saying why is this job's to do. + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + except Exception as e: + # A node failed, and the sink normally recorded which one. But the + # sink's own calls are outside the engine's try, so one of them + # failing -- a locked database, a run row deleted underneath -- is + # also how we get here, and then nothing was recorded at all. _fail + # only writes when the run is not already in a terminal state, so + # calling it either way cannot overwrite what the sink said. + self._fail(session_factory, pipeline_run_id, str(e)) + raise JobError(str(e)) from e + finally: + gc.collect() + + # The contexts of the leaves are all that is still held. Their live + # objects are datasets and models, which nothing will read again. + for ctx in contexts.values(): + ctx.clear_cache() + gc.collect() + + self.report_progress(1.0, "Finished") + + @staticmethod + def _fail(session_factory, pipeline_run_id: int, message: str) -> None: + """Mark the run as failed, unless something already settled it. + + Never overwrites a terminal status. The sink records a node failure + with the message the unit produced, which is the better one; this is + the backstop for the paths where nothing recorded anything, and a run + left in STARTED with no error is indistinguishable from one still + going. """ - Update the pipeline context and database object based on node type. - - Args: - context: The pipeline context dictionary - pipeline: The pipeline database object - node_type: The type of node that was executed - node_id: The ID of the node that was executed - output: The output from the node execution - """ - if node_type == "DataSelector": - context["dataset"] = output.get("dataset") - elif node_type == "DataExploration": - context["exploration"] = output.get("exploration") - pipeline.exploration = context["exploration"] - elif node_type == "Train": - context["train"] = output.get("train") - pipeline.train = context["train"] - elif node_type == "RetrieveModel": - context["retrieve"] = output.get("retrieve") - elif node_type == "Prediction": - context["prediction"] = output.get("prediction") - pipeline.prediction = context["prediction"] - else: - context[node_id] = output + from DashAI.back.core.enums.status import PipelineRunStatus + from DashAI.back.dependencies.database.models import PipelineRun + + settled = { + PipelineRunStatus.FINISHED, + PipelineRunStatus.ERROR, + } + + with session_factory() as db: + pipeline_run = db.get(PipelineRun, pipeline_run_id) + if pipeline_run is None or pipeline_run.status in settled: + return + pipeline_run.set_status_as_error(message) + db.commit() diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index 1d051a68a..35e09a1e3 100644 --- a/DashAI/back/models/base_model.py +++ b/DashAI/back/models/base_model.py @@ -224,55 +224,56 @@ def _save_metrics( db.commit() @final - def calculate_metrics( + def compute_metrics( self, split: SplitEnum = SplitEnum.VALIDATION, - level: LevelEnum = LevelEnum.LAST, - log_index: int = None, x_data: "DashAIDataset" = None, y_data: "DashAIDataset" = None, - ): - """Calculate and save metrics for a given data split and level. + ) -> Dict[str, float]: + """Score a data split with this model's metrics, without persisting. + + Which metrics are computed is decided by the model rather than by the + caller: ``ModelFactory`` attaches the metric classes and the data + splits to the instance, and this reads them off it. + + Separate from :meth:`calculate_metrics` so a caller with no run to log + against can still have the numbers. A pipeline is that caller: it has + no ``Run`` row, so there is no foreign key for a ``Metric`` row to + point at, and its results are recorded as an artifact instead. Both + paths score through here, so the two cannot disagree. Parameters ---------- split : SplitEnum - The data split to evaluate (TRAIN, VALIDATION, - or TEST). Defaults to SplitEnum.VALIDATION. - level : LevelEnum - The metric granularity level (LAST, TRIAL, - STEP, or BATCH). Defaults to LevelEnum.LAST. - log_index : int, optional - Explicit step index for the metric - entry. If None, the next step index is computed automatically. - Defaults to None. + The data split to evaluate. Defaults to ``SplitEnum.VALIDATION``. x_data : DashAIDataset, optional - Input features. If None, the - dataset stored in the model for the given split is used. - Defaults to None. + Input features. Defaults to the split stored on the model. y_data : DashAIDataset, optional - Target labels. If None, the - labels stored in the model for the given split are used. - Defaults to None. + Target labels. Defaults to the split stored on the model. + + Returns + ------- + Dict[str, float] + Metric name to score. Empty when every metric returned a + non-finite value. ``None`` when there is nothing to score at all: + no metrics configured, or no data for this split. """ - # Get the appropriate metrics based on split metrics_attr = f"{split.value}_metrics" metrics = getattr(self, metrics_attr, None) - # If no metrics or run_id, skip calculation - if not metrics or not self.run_id: - return + if not metrics: + return None # Load data if not provided if x_data is None or y_data is None: if self.x_data is None or self.y_data is None: - return + return None x_data = self.x_data[split.value] y_data = self.y_data[split.value] # If data is empty after retrieval, skip calculation if x_data is None or y_data is None: - return + return None # Make predictions and transform outputs y_pred = self.predict(x_data) @@ -293,6 +294,63 @@ def calculate_metrics( continue results[metric.__name__] = score + return results + + @final + def calculate_metrics( + self, + split: SplitEnum = SplitEnum.VALIDATION, + level: LevelEnum = LevelEnum.LAST, + log_index: int = None, + x_data: "DashAIDataset" = None, + y_data: "DashAIDataset" = None, + ): + """Calculate metrics for a data split and save them to the database. + + Parameters + ---------- + split : SplitEnum + The data split to evaluate (TRAIN, VALIDATION, + or TEST). Defaults to SplitEnum.VALIDATION. + level : LevelEnum + The metric granularity level (LAST, TRIAL, + STEP, or BATCH). Defaults to LevelEnum.LAST. + log_index : int, optional + Explicit step index for the metric + entry. If None, the next step index is computed automatically. + Defaults to None. + x_data : DashAIDataset, optional + Input features. If None, the + dataset stored in the model for the given split is used. + Defaults to None. + y_data : DashAIDataset, optional + Target labels. If None, the + labels stored in the model for the given split are used. + Defaults to None. + + Notes + ----- + A metric row is keyed by the run it belongs to, so a model with no run + has nowhere to write and this returns without scoring anything. That is + what lets a caller with no ``Run`` row -- a pipeline -- train a model + that logs nothing at all, during training or after it, without having + to intercept anything. Such a caller uses :meth:`compute_metrics` and + keeps the numbers itself. + """ + # No run means no foreign key for a metric row to point at. + # + # getattr rather than self.run_id: ModelFactory attaches the attribute, + # but a model built directly never had it, and the guard this replaced + # happened to never reach the attribute for such a model because it + # checked for metrics first. Treating "no attribute" as "no run" keeps + # that path working and is the same answer for any caller that has one. + if not getattr(self, "run_id", None): + return + + results = self.compute_metrics(split=split, x_data=x_data, y_data=y_data) + if results is None: + return + # Save to database self._save_metrics( split=split, level=level, results=results, log_index=log_index diff --git a/DashAI/back/optimizers/base_optimizer.py b/DashAI/back/optimizers/base_optimizer.py index 9d9750d6a..24bcb81c4 100644 --- a/DashAI/back/optimizers/base_optimizer.py +++ b/DashAI/back/optimizers/base_optimizer.py @@ -381,7 +381,7 @@ def importance_plot(self, trials, goal_metric): return PlotlyArtifact(payload=fig, title=title) - def create_plots(self, trials, run_id, n_params, goal_metric): + def create_plots(self, trials, run_id, n_params, goal_metric, artifact_prefix=None): """ List of available plots. @@ -393,18 +393,31 @@ def create_plots(self, trials, run_id, n_params, goal_metric): n_params (int): Number of the different hyperparameters involved in the process of hyperparameter optimization goal_metric (dict): Metric optimized in the process. + artifact_prefix (str, optional): Name to build the filenames from, + for callers that are not a run. Defaults to None, + which keeps deriving them from ``run_id``. Returns ------- plots_filenames (list): Filenames to persist each plot under. plots_list (list): The matching list of PlotlyArtifact instances. + + Notes + ----- + The filenames land in a directory shared with every other run's + artifacts, so whatever names them has to be unique per caller. + ``run_id`` is unique among runs but a pipeline has no run: it passes + ``artifact_prefix`` instead, and two pipeline executions that would + otherwise both write ``..._None.pickle`` stay apart. """ + tag = artifact_prefix if artifact_prefix is not None else run_id + if n_params >= 2: plots_filenames = [ - f"history_objective_plot_{run_id}.pickle", - f"slice_plot_{run_id}.pickle", - f"contour_plot_{run_id}.pickle", - f"importance_plot_{run_id}.pickle", + f"history_objective_plot_{tag}.pickle", + f"slice_plot_{tag}.pickle", + f"contour_plot_{tag}.pickle", + f"importance_plot_{tag}.pickle", ] plots_list = [ self.history_objective_plot(trials, goal_metric), @@ -415,8 +428,8 @@ def create_plots(self, trials, run_id, n_params, goal_metric): return plots_filenames, plots_list else: plots_filenames = [ - f"history_objective_plot_{run_id}.pickle", - f"slice_plot_{run_id}.pickle", + f"history_objective_plot_{tag}.pickle", + f"slice_plot_{tag}.pickle", ] plots_list = [ self.history_objective_plot(trials, goal_metric), diff --git a/DashAI/back/units/base_unit.py b/DashAI/back/units/base_unit.py index 3bb5c2947..351f0fc3d 100644 --- a/DashAI/back/units/base_unit.py +++ b/DashAI/back/units/base_unit.py @@ -37,7 +37,24 @@ class BaseUnit(ConfigObject, metaclass=ABCMeta): #: Context keys the unit guarantees after it runs. PROVIDES: Tuple[str, ...] = () + #: Configuration a user fills in. This is what the front renders. SCHEMA: BaseSchema = BaseSchema + #: Configuration names supplied by whatever runs the unit, and never by a + #: user. Deliberately kept out of ``SCHEMA``: a value here is not known when + #: a form would be filled in, so there is nothing for a user to answer. + #: ``artifact_prefix`` is the clearest case — it is built from the id of a + #: run that has not started. + #: + #: Being a separate declaration rather than a flag inside the schema is the + #: whole point. A flag would still travel to the front, where every renderer + #: would have to remember to skip it and a new one would leak by default. + #: These names never reach the schema the front receives, so there is + #: nothing to filter and nothing to forget. + #: + #: They still arrive through ``**config`` and are read as ``self.config[…]`` + #: like any other, so declaring one here changes nothing about how a unit is + #: written — only about who is expected to supply it. + RUNTIME_PARAMS: Tuple[str, ...] = () def __init__(self, **config) -> None: """Store the unit configuration. diff --git a/DashAI/back/units/build_manual_input_unit.py b/DashAI/back/units/build_manual_input_unit.py index 6cf31f19c..4a0c3ac71 100644 --- a/DashAI/back/units/build_manual_input_unit.py +++ b/DashAI/back/units/build_manual_input_unit.py @@ -35,31 +35,6 @@ class BuildManualInputSchema(BaseSchema): 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=[], @@ -102,6 +77,7 @@ class BuildManualInputUnit(BaseUnit): SCHEMA = BuildManualInputSchema PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("train_dataset_file_path",) def __init__(self, **config) -> None: super().__init__(**config) diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py index 3f39bb1b6..72336ad79 100644 --- a/DashAI/back/units/build_model_unit.py +++ b/DashAI/back/units/build_model_unit.py @@ -120,12 +120,18 @@ class BuildModelUnit(BaseUnit): SCHEMA = BuildModelSchema - # run_id and task_name only appear in the ModelFactory call and in error - # messages, but they are declared all the same: a key read without being - # declared is invisible to any caller — and to any future DAG validator — - # that inspects REQUIRES instead of running the unit. - REQUIRES = ("x", "y", "n_labels", "run_id", "task_name") + # task_name only appears in the ModelFactory call and in error messages, + # but it is declared all the same: a key read without being declared is + # invisible to any caller — and to the DAG validator — that inspects + # REQUIRES instead of running the unit. + # + # run_id is configuration, not context: no unit publishes it, so nothing + # upstream could ever satisfy it. It is read without a default on purpose + # — a run_id nobody passed would read as "this model has no run", and a + # model with no run logs no metrics at all (see BaseModel). + REQUIRES = ("x", "y", "n_labels", "task_name") PROVIDES = ("model", "factory", "optimizable_parameters", "model_parameters") + RUNTIME_PARAMS = ("run_id",) def __init__(self, **config) -> None: super().__init__(**config) @@ -199,7 +205,7 @@ def execute(self, ctx: ExecutionContext) -> None: component_registry = di["component_registry"] parameters = self.model_parameters - run_id = ctx.require("run_id") + run_id = self.config["run_id"] task_name = ctx.require("task_name") model_class = self._resolve_model_class() diff --git a/DashAI/back/units/compute_dataset_metadata_unit.py b/DashAI/back/units/compute_dataset_metadata_unit.py index d903b5263..124c28e5b 100644 --- a/DashAI/back/units/compute_dataset_metadata_unit.py +++ b/DashAI/back/units/compute_dataset_metadata_unit.py @@ -58,35 +58,6 @@ class ComputeDatasetMetadataSchema(BaseSchema): zh="计算扩展元数据", ), ) # type: ignore - trust_inherited_metadata: schema_field( - bool_field(), - placeholder=False, - description=MultilingualString( - en="Whether the metadata the dataset already carries can be reused " - "as-is. Only safe when the data has not changed since that metadata " - "was computed; otherwise every value is recomputed from scratch.", - es="Si la metadata que el conjunto de datos ya trae puede reusarse " - "tal cual. Solo es seguro cuando los datos no cambiaron desde que " - "esa metadata se calculó; si no, todo se recalcula desde cero.", - pt="Se a metadata que o conjunto de dados já carrega pode ser " - "reutilizada como está. Só é seguro quando os dados não mudaram " - "desde que essa metadata foi calculada; caso contrário, tudo é " - "recalculado do zero.", - de="Ob die vom Datensatz mitgeführten Metadaten unverändert " - "weiterverwendet werden können. Nur sicher, wenn sich die Daten " - "seit deren Berechnung nicht geändert haben; andernfalls wird alles " - "neu berechnet.", - zh="数据集已携带的元数据是否可以原样重用。仅当数据自该元数据计算后" - "未发生变化时才安全；否则将全部重新计算。", - ), - alias=MultilingualString( - en="Trust existing metadata", - es="Confiar en la metadata existente", - pt="Confiar na metadata existente", - de="Vorhandene Metadaten vertrauen", - zh="信任现有元数据", - ), - ) # type: ignore class ComputeDatasetMetadataUnit(BaseUnit): @@ -113,6 +84,7 @@ class ComputeDatasetMetadataUnit(BaseUnit): REQUIRES = ("dataset",) PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("trust_inherited_metadata",) def execute(self, ctx: ExecutionContext) -> None: dataset = ctx.require("dataset") diff --git a/DashAI/back/units/context.py b/DashAI/back/units/context.py index 7782ec93f..daa1ef6a0 100644 --- a/DashAI/back/units/context.py +++ b/DashAI/back/units/context.py @@ -153,6 +153,36 @@ def has(self, key: str) -> bool: """Return whether a key is present in either half of the context.""" return key in self._cache or key in self._refs + def origin(self, key: str) -> Optional[str]: + """Return which half holds a key, without copying anything. + + ``get`` and ``has`` merge the two halves on purpose, and ``refs`` + deep-copies the whole reference half, so neither can answer this + cheaply. Something that moves a value from one context to another has + to know: the two halves have incompatible rules, and picking the wrong + one fails in both directions. Handing a live dataset to ``put_ref`` + raises, and handing ``put`` a dict that was a reference silently drops + the copy-on-read guarantee the reference depended on. + + The cache is checked first, matching ``get`` and ``require``, so a key + present in both halves has one answer rather than two. + + Parameters + ---------- + key : str + Name of the value. + + Returns + ------- + Optional[str] + ``"cache"``, ``"ref"``, or ``None`` when the key is absent. + """ + if key in self._cache: + return "cache" + if key in self._refs: + return "ref" + return None + def clear_cache(self) -> None: """Drop every live object, keeping the references. diff --git a/DashAI/back/units/evaluate_model_to_artifact_unit.py b/DashAI/back/units/evaluate_model_to_artifact_unit.py new file mode 100644 index 000000000..14e73ceb1 --- /dev/null +++ b/DashAI/back/units/evaluate_model_to_artifact_unit.py @@ -0,0 +1,90 @@ +"""Unit that computes a trained model's metrics and publishes them.""" + +import logging + +from DashAI.back.core.enums.metrics import SplitEnum +from DashAI.back.core.schema_fields import ( + BaseSchema, + enum_field, + list_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 + +log = logging.getLogger(__name__) + +DEFAULT_SPLITS = ["TRAIN", "VALIDATION", "TEST"] + + +class EvaluateModelToArtifactSchema(BaseSchema): + splits: schema_field( + list_field(enum_field(enum=DEFAULT_SPLITS)), + placeholder=DEFAULT_SPLITS, + description=MultilingualString( + en="Data splits the model is evaluated on.", + es="Particiones de datos sobre las que se evalúa el modelo.", + pt="Partições de dados sobre as quais o modelo é avaliado.", + de="Datenteilmengen, auf denen das Modell ausgewertet wird.", + zh="用于评估模型的数据划分。", + ), + alias=MultilingualString( + en="Splits", + es="Particiones", + pt="Partições", + de="Teilmengen", + zh="数据划分", + ), + ) # type: ignore + + +class EvaluateModelToArtifactUnit(BaseUnit): + """Compute a trained model's metrics and publish them as data. + + The sibling of ``EvaluateModelUnit``, for a caller that has no ``Run`` row. + That unit writes ``Metric`` rows against a foreign key to ``run.id``, so it + cannot work without one; this returns the same numbers instead of + persisting them, and whoever asked for them decides where they go. + + Which metrics are computed is decided by the model, not by this unit, and + not configured here either: ``ModelFactory`` attaches the metric classes + and the data splits to the instance, so the metrics stay configured in one + place — where the model is built — rather than in two nodes that could + disagree. + + Both this and ``EvaluateModelUnit`` score through + ``BaseModel.compute_metrics``, which is what keeps them from drifting apart + into two answers for the same model and split. + """ + + SCHEMA = EvaluateModelToArtifactSchema + + REQUIRES = ("model",) + PROVIDES = ("metrics",) + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + splits = [SplitEnum[name] for name in self.config.get("splits", DEFAULT_SPLITS)] + + metrics = {} + try: + for split in splits: + scores = model.compute_metrics(split=split) + # None means there was nothing to score -- no metrics + # configured, or no data for this split. Recording an empty + # entry would claim the split was evaluated and scored zero + # metrics, which is a different statement. + if scores is None: + continue + metrics[split.value] = scores + except Exception as e: + log.exception(e) + raise JobError( + f"Metric calculation failed {e}", + ) from e + + # A ref, not a cached object: plain numbers are exactly what survives + # leaving the process, and what a caller records. + ctx.put_ref("metrics", metrics) diff --git a/DashAI/back/units/evaluate_model_unit.py b/DashAI/back/units/evaluate_model_unit.py index bdce525f7..edcfa0c9e 100644 --- a/DashAI/back/units/evaluate_model_unit.py +++ b/DashAI/back/units/evaluate_model_unit.py @@ -52,11 +52,33 @@ class EvaluateModelUnit(BaseUnit): ``BaseModel.calculate_metrics`` is ``final``. That method persists the rows through a session of its own, so these writes are not part of the transaction the calling job controls. + + **This unit needs a real ``Run`` row and cannot be used without one.** + Everything it does is write ``Metric`` rows against a foreign key to + ``run.id``, so there is nothing left for it to do when there is no run. + ``validate`` refuses a missing run id rather than letting it through, + because the failure would otherwise be silent in both directions: the + idempotency query below would match no row whatever was already logged, + and ``calculate_metrics`` no-ops on a model with no run — so the unit + would report success having written nothing at all. + + A caller that wants a model's metrics *without* a run wants a different + unit, one that returns them instead of persisting them. """ SCHEMA = EvaluateModelSchema - REQUIRES = ("model", "run_id") + # run_id is configuration, not context: no unit publishes it. + REQUIRES = ("model",) + RUNTIME_PARAMS = ("run_id",) + + def validate(self, ctx: ExecutionContext) -> None: + if self.config["run_id"] is None: + raise JobError( + "Metrics can only be logged against a run, and this one has no " + "run id. Use a unit that returns the metrics instead of writing " + "them if there is no run to attach them to." + ) def execute(self, ctx: ExecutionContext) -> None: from kink import di @@ -64,13 +86,10 @@ def execute(self, ctx: ExecutionContext) -> None: session_factory = di["session_factory"] model = ctx.require("model") - # ctx.require, not ctx.get: run_id is what the idempotency query below - # filters on. A silently-None run_id would match no existing metric - # row regardless of what was actually logged, and — if the model - # were also somehow detached from its run — calculate_metrics would - # then no-op (base_model.py's ``if not metrics or not self.run_id``), - # so the unit would "succeed" having written nothing. - run_id = ctx.require("run_id") + # validate() already refused a None here, which is what keeps the + # idempotency query below from matching no row regardless of what was + # actually logged. + run_id = self.config["run_id"] splits = [SplitEnum[name] for name in self.config.get("splits", DEFAULT_SPLITS)] try: diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 8f9c6bf14..49f34add3 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -84,6 +84,9 @@ class FitModelUnit(BaseUnit): SCHEMA = FitModelSchema + # run_id and artifact_prefix are configuration, not context: no unit + # publishes them, so nothing upstream could ever satisfy them as REQUIRES. + # See the artifact_prefix section of DAG_ENGINE.md. REQUIRES = ( "model", "factory", @@ -92,9 +95,9 @@ class FitModelUnit(BaseUnit): "x", "y", "task", - "run_id", ) PROVIDES = ("model", "plot_paths") + RUNTIME_PARAMS = ("run_id", "artifact_prefix") def __init__(self, **config) -> None: super().__init__(**config) @@ -167,7 +170,7 @@ def execute(self, ctx: ExecutionContext) -> None: model = ctx.require("model") x = ctx.require("x") y = ctx.require("y") - run_id = ctx.require("run_id") + run_id = self.config["run_id"] optimizable_parameters = ctx.require("optimizable_parameters") plot_paths = [] @@ -191,7 +194,7 @@ def execute(self, ctx: ExecutionContext) -> None: model = optimizer.get_model() best_params = optimizer.get_best_params() - self._assert_model_keeps_its_runtime_state(model, run_id) + self._assert_model_keeps_its_runtime_state(model) # ctx.require already hands back an isolated copy of the # stored parameter tree, so update_parameters is free to @@ -211,6 +214,7 @@ def execute(self, ctx: ExecutionContext) -> None: run_id, n_params=len(optimizable_parameters), goal_metric=goal_metric, + artifact_prefix=self.config["artifact_prefix"], ) normalized_plots = normalize_artifacts(plots) for filename, plot in zip( @@ -230,21 +234,25 @@ def execute(self, ctx: ExecutionContext) -> None: ctx.put_ref("plot_paths", plot_paths) @staticmethod - def _assert_model_keeps_its_runtime_state(model, run_id) -> None: - """Fail loudly if the optimizer returned a model that cannot log metrics. - - ``ModelFactory`` attaches the run id, the data splits and the metric - classes to the model instance, and optimizers are expected to return - that same instance. If one ever returns a fresh object instead, - ``calculate_metrics`` would return early and the run would finish with - no metrics at all instead of failing. + def _assert_model_keeps_its_runtime_state(model) -> None: + """Fail loudly if the optimizer returned a model that cannot be scored. + + ``ModelFactory`` attaches the data splits and the metric classes to the + model instance, and optimizers are expected to return that same + instance. If one ever returns a fresh object instead, scoring it finds + nothing to score and the caller ends up with no metrics rather than an + error. + + The check is on the data, not on the run id. Keying it to ``run_id`` + made it a no-op for every caller that has no run -- a pipeline, where + ``run_id`` is always None -- which is exactly the caller with no other + signal that anything went wrong: it would finish with an empty metrics + artifact. What both callers need is the same, so this asks for that + instead. """ - if run_id is None: - return - - if getattr(model, "run_id", None) is None: + if getattr(model, "x_data", None) is None: raise JobError( - "The optimizer returned a model detached from its run: metrics " + "The optimizer returned a model detached from its data: metrics " "could not be computed for it. Optimizers must return the same " "model instance they received." ) diff --git a/DashAI/back/units/generate_local_explanation_unit.py b/DashAI/back/units/generate_local_explanation_unit.py index 7f12803ab..32cca16fd 100644 --- a/DashAI/back/units/generate_local_explanation_unit.py +++ b/DashAI/back/units/generate_local_explanation_unit.py @@ -202,32 +202,6 @@ class GenerateLocalExplanationSchema(BaseSchema): 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): @@ -252,6 +226,7 @@ class GenerateLocalExplanationUnit(BaseUnit): REQUIRES = ("explainer", "data_x", "data_y", "task", "split_indexes") PROVIDES = ("explanation_path", "plots_path", "input_dataset_path") + RUNTIME_PARAMS = ("session_splits",) def _select_instances(self, prepared_instance, splits, instance, task): """Narrow the loaded dataset down to the instances to explain.""" diff --git a/DashAI/back/units/load_training_dataset_unit.py b/DashAI/back/units/load_training_dataset_unit.py index 0cdd69f75..8d92280c3 100644 --- a/DashAI/back/units/load_training_dataset_unit.py +++ b/DashAI/back/units/load_training_dataset_unit.py @@ -3,12 +3,6 @@ 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 @@ -19,34 +13,6 @@ 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. @@ -63,9 +29,8 @@ class LoadTrainingDatasetUnit(BaseUnit): from the dataset *being predicted on* crosses this boundary. """ - SCHEMA = LoadTrainingDatasetSchema - PROVIDES = ("train_dataset", "train_dataset_types") + RUNTIME_PARAMS = ("train_dataset_file_path",) def execute(self, ctx: ExecutionContext) -> None: from pathlib import Path diff --git a/DashAI/back/units/load_uploaded_dataset_unit.py b/DashAI/back/units/load_uploaded_dataset_unit.py index 0b708d0d2..a287b8f17 100644 --- a/DashAI/back/units/load_uploaded_dataset_unit.py +++ b/DashAI/back/units/load_uploaded_dataset_unit.py @@ -63,28 +63,6 @@ class LoadUploadedDatasetSchema(BaseSchema): zh="来源", ), ) # type: ignore - temp_path: schema_field( - string_field(), - placeholder="", - description=MultilingualString( - en="Scratch directory for downloads and extracted archives. Whoever " - "sets it up is responsible for removing it afterwards.", - es="Directorio temporal para descargas y archivos extraídos. Quien " - "lo crea es responsable de borrarlo después.", - pt="Diretório temporário para descargas e ficheiros extraídos. Quem " - "o cria é responsável por removê-lo depois.", - de="Arbeitsverzeichnis für Downloads und entpackte Archive. Wer es " - "anlegt, ist für das Entfernen verantwortlich.", - zh="用于下载和解压归档的临时目录。创建者负责事后清理。", - ), - alias=MultilingualString( - en="Temporary path", - es="Ruta temporal", - pt="Caminho temporário", - de="Temporärer Pfad", - zh="临时路径", - ), - ) # type: ignore n_sample: schema_field( none_type(int_field(gt=0)), placeholder=None, @@ -125,6 +103,7 @@ class LoadUploadedDatasetUnit(BaseUnit): SCHEMA = LoadUploadedDatasetSchema PROVIDES = ("dataset",) + RUNTIME_PARAMS = ("temp_path",) def execute(self, ctx: ExecutionContext) -> None: from kink import di diff --git a/DashAI/back/units/save_model_unit.py b/DashAI/back/units/save_model_unit.py index 22d629b12..fa808b61f 100644 --- a/DashAI/back/units/save_model_unit.py +++ b/DashAI/back/units/save_model_unit.py @@ -1,6 +1,7 @@ """Unit that persists a trained model to disk.""" import logging +import re from DashAI.back.job.base_job import JobError from DashAI.back.units.base_unit import BaseUnit @@ -8,16 +9,40 @@ log = logging.getLogger(__name__) +#: A prefix names a directory directly under RUNS_PATH, so anything that could +#: be read as a path — a separator, a parent reference, a drive letter — has to +#: be refused rather than cleaned up: a caller that meant one destination and +#: silently got another is the failure this guards against. +_SAFE_PREFIX = re.compile(r"^[A-Za-z0-9_-]+$") + class SaveModelUnit(BaseUnit): - """Write a trained model under the runs directory, keyed by its run id. + """Write a trained model under the runs directory, in its own subdirectory. + + The destination is named by ``artifact_prefix``, which the caller chooses: + a run passes its own id, so a re-run overwrites its own artifact and never + another's. A caller that is not a run passes something unique to itself. - Takes no configuration: the destination is derived from the run the model - belongs to, so a re-run overwrites its own artifact and never another's. + The prefix is configuration rather than a context key on purpose: no unit + publishes it, so nothing upstream could ever satisfy it as a requirement. + It is the same split ``SaveDatasetUnit`` and ``SaveDatasetToPathUnit`` + already draw — a destination an upstream unit produced against one the + caller names. """ - REQUIRES = ("model", "run_id") + REQUIRES = ("model",) PROVIDES = ("model_path",) + RUNTIME_PARAMS = ("artifact_prefix",) + + def validate(self, ctx: ExecutionContext) -> None: + prefix = self.config["artifact_prefix"] + + if not isinstance(prefix, str) or not _SAFE_PREFIX.match(prefix): + raise JobError( + "The artifact prefix names a directory under the runs " + "directory, so it can only contain letters, digits, hyphens " + f"and underscores. Got: {prefix!r}" + ) def execute(self, ctx: ExecutionContext) -> None: import os @@ -27,10 +52,11 @@ def execute(self, ctx: ExecutionContext) -> None: config = di["config"] model = ctx.require("model") - run_id = ctx.require("run_id") try: - model_path = os.path.join(config["RUNS_PATH"], str(run_id)) + model_path = os.path.join( + config["RUNS_PATH"], self.config["artifact_prefix"] + ) model.save(model_path) except Exception as e: log.exception(e) diff --git a/tests/back/api/test_model_job_as_a_graph.py b/tests/back/api/test_model_job_as_a_graph.py new file mode 100644 index 000000000..a8ce2473c --- /dev/null +++ b/tests/back/api/test_model_job_as_a_graph.py @@ -0,0 +1,535 @@ +"""ModelJob's whole shape, run as a graph. The test that says the design works. + +The spike that established the engine never ran this, and the three problems it +left open all live here: where run_id goes, whether the bundling rule survives +fifteen edges, and whether a training pipeline can produce metrics with no Run +row to hang them on. + +Six nodes, and the pipeline is a sandbox: it creates no Run and no ModelSession, +writes no Metric rows, and keeps its model out of the way of every real run's. + + load --dataset,dataset_id--> prep --x,y,n_labels,task_name--> build + | | + +-------x,y,task-------------> fit <-+ + | + +-------model-------+-------+ + v v + eval save +""" + +import json +import os + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dataloaders.classes.csv_dataloader import CSVDataLoader +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.job.base_job import JobError +from DashAI.back.job.pipeline_job import PipelineJob +from DashAI.back.metrics.base_metric import BaseMetric +from DashAI.back.models.base_model import BaseModel +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit + +SPLITS = { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + "splitType": "random", +} + + +class GraphTask(BaseTask): + name: str = "GraphTask" + 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 + + def num_labels(self, dataset, output_column): + return 3 + + +#: How many times a model asked to log metrics while training, across the run. +#: A model instance does not survive to be inspected -- the graph releases it -- +#: and the count is what makes "no metric rows" a consequence rather than a +#: coincidence. +LOG_ATTEMPTS = [] + + +class GraphModel(BaseModel): + """Logs metrics while training, the way the torch-based models do. + + That matters for what this file is testing. A model that never asks to log + would leave the Metric table empty no matter what the pipeline did, so the + sandbox would look airtight without being tested at all. This one asks on + every train call, so an empty Metric table means the switch stopped it. + """ + + COMPATIBLE_COMPONENTS = ["GraphTask"] + + def __init__(self, **kwargs): + self.trained_with = None + + def save(self, filename): + joblib.dump({"trained_with": self.trained_with}, filename) + + def load(self, filename): + return joblib.load(filename) + + def predict(self, x): + return [0] * x.shape[0] + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.trained_with = {"train": x_train.shape[0]} + LOG_ATTEMPTS.append(getattr(self, "run_id", "missing")) + self.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.STEP) + return self + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + def prepare_output(self, dataset, is_fit=False): + return [0] * len(LOG_ATTEMPTS or [0]) + + +class GraphMetric(BaseMetric): + COMPATIBLE_COMPONENTS = ["GraphTask"] + MAXIMIZE = True + + @staticmethod + def score(true_labels, probs_pred_labels): + return 0.5 + + +@pytest.fixture(scope="module", name="graph_registry", autouse=True) +def setup_graph_registry(client): + services = client.app.container._services + sentinel = object() + old = services.get("component_registry", sentinel) + + services["component_registry"] = ComponentRegistry( + initial_components=[ + GraphTask, + GraphModel, + GraphMetric, + CSVDataLoader, + OptunaOptimizer, + PipelineJob, + LoadDatasetUnit, + PrepareAndSplitUnit, + BuildModelUnit, + FitModelUnit, + EvaluateModelToArtifactUnit, + SaveModelUnit, + ] + ) + yield services["component_registry"] + if old is sentinel: + del services["component_registry"] + else: + services["component_registry"] = old + + +def _model_job_blocks(dataset_id: int): + """The six blocks, and the six drawn edges that stand for fifteen wires.""" + steps = [ + { + "id": "load", + "units": [ + { + "id": "load", + "unit": "LoadDatasetUnit", + "config": {"dataset_id": dataset_id}, + } + ], + }, + { + "id": "prep", + "units": [ + { + "id": "prep", + "unit": "PrepareAndSplitUnit", + "config": { + "task_name": "GraphTask", + "input_columns": ["SepalLengthCm", "SepalWidthCm"], + "output_columns": ["Species"], + "splits": SPLITS, + }, + } + ], + }, + { + "id": "build", + "units": [ + { + "id": "build", + "unit": "BuildModelUnit", + "config": { + "model": {"component": "GraphModel", "params": {}}, + "train_metrics": ["GraphMetric"], + "validation_metrics": ["GraphMetric"], + "test_metrics": ["GraphMetric"], + }, + } + ], + }, + { + "id": "fit", + "units": [ + { + "id": "fit", + "unit": "FitModelUnit", + "config": { + "optimizer": { + "component": "OptunaOptimizer", + "params": {}, + }, + "goal_metric": "GraphMetric", + }, + } + ], + }, + { + "id": "eval", + "units": [ + { + "id": "eval", + "unit": "EvaluateModelToArtifactUnit", + "config": {"splits": ["TRAIN", "VALIDATION", "TEST"]}, + } + ], + }, + { + "id": "save", + "units": [{"id": "save", "unit": "SaveModelUnit", "config": {}}], + }, + ] + edges = [ + {"source": "load", "target": "prep"}, + {"source": "prep", "target": "build"}, + {"source": "prep", "target": "fit"}, + {"source": "build", "target": "fit"}, + {"source": "fit", "target": "eval"}, + {"source": "fit", "target": "save"}, + ] + return steps, edges + + +@pytest.fixture(name="finished_pipeline_run", scope="module") +def run_the_graph(client: TestClient, dataset_1: Dataset, graph_registry): + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + with session_factory() as db: + pipeline = Pipeline(name="ModelJob as a graph", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + pipeline_run = db.query(PipelineRun).filter_by(pipeline_id=pipeline_id).one() + db.refresh(pipeline_run) + # Read everything now: the session closes with the fixture. + yield { + "id": pipeline_run.id, + "status": pipeline_run.status, + "error_message": pipeline_run.error_message, + "steps": pipeline_run.steps, + "edges": pipeline_run.edges, + "nodes": { + row.node_id: { + "status": row.status, + "start_time": row.start_time, + "end_time": row.end_time, + "block_id": row.block_id, + "node_type": row.node_type, + "config": row.config, + } + for row in pipeline_run.node_runs + }, + "artifacts": { + (row.node_run.node_id, row.key): row.value + for row in db.query(NodeArtifact).all() + }, + } + + +def test_the_whole_graph_runs_to_completion(finished_pipeline_run): + assert finished_pipeline_run["status"] == PipelineRunStatus.FINISHED + assert finished_pipeline_run["error_message"] is None + + nodes = finished_pipeline_run["nodes"] + assert set(nodes) == {"load", "prep", "build", "fit", "eval", "save"} + for node_id, node in nodes.items(): + assert node["status"] == NodeRunStatus.FINISHED, node_id + assert node["start_time"] is not None, node_id + assert node["end_time"] is not None, node_id + + +def test_six_drawn_edges_expand_into_fifteen_wires(finished_pipeline_run): + """The granularity problem, measured. + + The unit contract is finer than a canvas can draw: FitModelUnit alone + requires seven keys. One drawn edge carries every key its two units agree + on, which is what makes six shapes on a canvas enough for this graph. + """ + edges = finished_pipeline_run["edges"] + assert len(edges) == 15 + + carried = {} + for edge in edges: + carried.setdefault((edge["src"], edge["dst"]), set()).add(edge["src_key"]) + + assert carried == { + ("load", "prep"): {"dataset", "dataset_id"}, + ("prep", "build"): {"x", "y", "n_labels", "task_name"}, + ("prep", "fit"): {"x", "y", "task"}, + ("build", "fit"): { + "model", + "factory", + "optimizable_parameters", + "model_parameters", + }, + ("fit", "eval"): {"model"}, + ("fit", "save"): {"model"}, + } + + +def test_the_sandbox_writes_no_metric_rows(client, finished_pipeline_run): + """The assertion that decides whether the sandbox holds. + + Metric.run_id is a foreign key to run.id, and SQLite is not enforcing it + here -- there is no PRAGMA foreign_keys=ON anywhere. So a row written with + a run id that matches nothing inserts happily, and pipeline_run.id and + run.id are independent sequences that both start at 1: a pipeline metric + would land in the metric list of the real run with that id and show up in + its live chart. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.query(Metric).count() == 0 + + +def test_the_model_did_ask_to_log_and_was_stopped(finished_pipeline_run): + """What makes the empty Metric table a consequence, not a coincidence. + + GraphModel calls calculate_metrics on every train call, the way the + torch-based models do. So the run above genuinely tried to write metric + rows and wrote none, and it tried with a model that had no run -- which is + the only thing standing between it and a row keyed by nothing. + """ + assert LOG_ATTEMPTS, "the model never tried to log, so nothing was tested" + assert set(LOG_ATTEMPTS) == {None} + + +def test_the_sandbox_creates_no_run_and_no_model_session(client, finished_pipeline_run): + """No application entity is manufactured to make a pipeline work. + + A Run needs a ModelSession, which needs a globally unique name, a dataset, + a task, input and output columns and splits, all NOT NULL -- two entities + per execution, visible in the Models UI as sessions nobody created, with a + deletion lifecycle nobody owns. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + assert db.query(Run).count() == 0 + assert db.query(ModelSession).count() == 0 + + +def test_the_metrics_of_the_pipeline_live_in_an_artifact(finished_pipeline_run): + """Where a pipeline's metrics go, now that they cannot go to Metric.""" + metrics = finished_pipeline_run["artifacts"][("eval", "metrics")] + + assert metrics == { + "train": {"GraphMetric": 0.5}, + "validation": {"GraphMetric": 0.5}, + "test": {"GraphMetric": 0.5}, + } + + +def test_the_model_is_saved_where_no_real_run_could_collide_with_it( + client, finished_pipeline_run +): + """The artifact name, which is the other half of what run_id used to do.""" + model_path = finished_pipeline_run["artifacts"][("save", "model_path")] + runs_path = str(client.app.container["config"]["RUNS_PATH"]) + + assert model_path.startswith(runs_path) + assert os.path.exists(model_path) + + directory = os.path.basename(model_path) + # Not a bare integer, which is what every real run's directory is called. + assert not directory.isdigit() + assert directory == f"pipeline-{finished_pipeline_run['id']}-save" + + # And the model that landed there is the one that was trained. + assert joblib.load(model_path)["trained_with"]["train"] > 0 + + +def test_the_model_never_belonged_to_a_run(finished_pipeline_run): + """The switch that turns metric persistence off, as it reaches the node. + + ModelFactory hangs run_id on the model instance, and + BaseModel.calculate_metrics returns early for a model with no run. That is + what makes the sandbox airtight without the pipeline having to intercept + anything: not even a model that logs while training can write a row. + """ + assert finished_pipeline_run["nodes"]["build"]["config"]["run_id"] is None + assert finished_pipeline_run["nodes"]["fit"]["config"]["run_id"] is None + + +def test_the_run_records_units_not_canvas_shapes(finished_pipeline_run): + nodes = finished_pipeline_run["nodes"] + + assert nodes["build"]["node_type"] == "BuildModelUnit" + assert nodes["eval"]["node_type"] == "EvaluateModelToArtifactUnit" + # One unit per block for now, so each node is its own block. + for node_id, node in nodes.items(): + assert node["block_id"] == node_id + + +def test_every_serializable_output_is_recorded(finished_pipeline_run): + """Artifacts are named by the unit's own PROVIDES keys.""" + artifacts = finished_pipeline_run["artifacts"] + + assert ("load", "dataset_id") in artifacts + assert ("load", "dataset_path") in artifacts + assert ("prep", "split_indexes") in artifacts + assert ("prep", "task_name") in artifacts + assert ("build", "model_parameters") in artifacts + assert ("fit", "plot_paths") in artifacts + assert ("eval", "metrics") in artifacts + assert ("save", "model_path") in artifacts + + # The live objects are not: x, y, the task, the model and the factory are + # all cache-half values, derivable again and never serialized. + for key in ("dataset", "x", "y", "task", "model", "factory"): + assert not any(recorded == key for _, recorded in artifacts) + + +def test_a_node_failing_halfway_cancels_the_rest( + client: TestClient, dataset_1: Dataset, graph_registry +): + """A run that died halfway must not look like one still in flight.""" + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + # A model that is not in the registry: build fails, and everything the + # order puts after it never runs. + for step in steps: + if step["id"] == "build": + step["units"][0]["config"]["model"] = { + "component": "NoSuchModel", + "params": {}, + } + + with session_factory() as db: + pipeline = Pipeline(name="A graph that fails", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + with pytest.raises(JobError): + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + pipeline_run = db.query(PipelineRun).filter_by(pipeline_id=pipeline_id).one() + statuses = {row.node_id: row.status for row in pipeline_run.node_runs} + assert pipeline_run.status == PipelineRunStatus.ERROR + + assert statuses["load"] == NodeRunStatus.FINISHED + assert statuses["prep"] == NodeRunStatus.FINISHED + assert statuses["build"] == NodeRunStatus.ERROR + assert statuses["fit"] == NodeRunStatus.CANCELLED + assert statuses["eval"] == NodeRunStatus.CANCELLED + assert statuses["save"] == NodeRunStatus.CANCELLED + + +def test_two_runs_of_the_same_graph_do_not_overwrite_each_others_model( + client: TestClient, dataset_1: Dataset, graph_registry +): + """Each execution names its own artifacts. + + Without the run id in the name, the second execution would write over the + first one's model, silently. + """ + session_factory = client.app.container["session_factory"] + steps, edges = _model_job_blocks(dataset_1.id) + + with session_factory() as db: + pipeline = Pipeline(name="A graph run twice", steps=steps, edges=edges) + db.add(pipeline) + db.commit() + pipeline_id = pipeline.id + + PipelineJob(pipeline_id=pipeline_id).run() + PipelineJob(pipeline_id=pipeline_id).run() + + with session_factory() as db: + runs = ( + db.query(PipelineRun) + .filter_by(pipeline_id=pipeline_id) + .order_by(PipelineRun.id) + .all() + ) + paths = [ + db.query(NodeArtifact) + .join(NodeRun) + .filter( + NodeRun.pipeline_run_id == pipeline_run.id, + NodeArtifact.key == "model_path", + ) + .one() + .value + for pipeline_run in runs + ] + + assert len(paths) == 2 + assert paths[0] != paths[1] + assert all(os.path.exists(path) for path in paths) + + +def test_the_graph_definition_survives_being_stored_as_json( + client: TestClient, dataset_1: Dataset, graph_registry +): + """The blocks are plain data, which is what a canvas can save and reload.""" + steps, edges = _model_job_blocks(dataset_1.id) + + assert json.loads(json.dumps(steps)) == steps + assert json.loads(json.dumps(edges)) == edges diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index dd5726b8d..c965eb3d9 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -9,6 +9,7 @@ "BuildModelUnit", "FitModelUnit", "EvaluateModelUnit", + "EvaluateModelToArtifactUnit", "SaveModelUnit", "ApplyConverterUnit", "FitConverterUnit", @@ -114,10 +115,11 @@ def test_unit_schemas_describe_their_configuration(units): assert set(units["GenerateGlobalExplanationUnit"]["schema"]["properties"]) == { "explainer_id" } + # temp_path is a runtime param: the endpoint makes the directory, so it + # never reaches the schema the front receives. assert set(units["LoadUploadedDatasetUnit"]["schema"]["properties"]) == { "dataloader", "source", - "temp_path", "n_sample", } assert set(units["LoadDatafileDatasetUnit"]["schema"]["properties"]) == { @@ -133,10 +135,11 @@ def test_unit_schemas_describe_their_configuration(units): } assert set(units["ComputeDatasetMetadataUnit"]["schema"]["properties"]) == { "compute_metadata", - "trust_inherited_metadata", } # The sibling of SaveDatasetUnit: that one saves where the load said, this - # one is told where to save. + # one is told where to save. The destination stays user-facing: a job + # happens to compute it, but on a canvas picking where to store a dataset + # is a decision, and a value a user can know while filling the form in. assert set(units["SaveDatasetToPathUnit"]["schema"]["properties"]) == {"path"} @@ -193,3 +196,70 @@ def test_units_do_not_leak_into_the_job_listing(client: TestClient): job_names = {component["name"] for component in response.json()} assert not (job_names & EXPECTED_UNITS) assert "ModelJob" in job_names + + +@pytest.fixture(name="unit_classes", scope="module") +def get_unit_classes(client: TestClient): + """The registered classes behind the response, from the same registry.""" + registry = client.app.container["component_registry"] + return {name: registry[name]["class"] for name in EXPECTED_UNITS} + + +#: Every runtime param in the palette. Configuration supplied by whatever runs +#: the unit -- a job, the DAG engine, an endpoint -- and never by a user, +#: because at the moment a form would be filled in the value is not knowable: +#: a path a job has yet to choose, a foreign key to a row nobody has seen, or +#: a name built from the id of a run that has not started. +EXPECTED_RUNTIME_PARAMS = { + ("LoadUploadedDatasetUnit", "temp_path"), + ("LoadTrainingDatasetUnit", "train_dataset_file_path"), + ("BuildManualInputUnit", "train_dataset_file_path"), + ("GenerateLocalExplanationUnit", "session_splits"), + ("ComputeDatasetMetadataUnit", "trust_inherited_metadata"), + ("BuildModelUnit", "run_id"), + ("EvaluateModelUnit", "run_id"), + ("FitModelUnit", "run_id"), + ("FitModelUnit", "artifact_prefix"), + ("SaveModelUnit", "artifact_prefix"), +} + + +def test_the_runtime_params_are_exactly_the_declared_ones(unit_classes): + """Adding or removing one has to be a deliberate edit, not a side effect.""" + declared = { + (name, param) + for name, cls in unit_classes.items() + for param in cls.RUNTIME_PARAMS + } + + assert declared == EXPECTED_RUNTIME_PARAMS + + +def test_no_runtime_param_reaches_the_schema_the_front_receives(units, unit_classes): + """The guarantee the whole arrangement exists for. + + A flag inside the schema would have left these on the wire, where every + renderer has to remember to skip them and a new one leaks by default. Being + a separate declaration means there is nothing to filter: the name is simply + not in what the front is given. + """ + for name, cls in unit_classes.items(): + exposed = set(units[name]["schema"]["properties"]) + leaked = exposed & set(cls.RUNTIME_PARAMS) + + assert not leaked, ( + f"{name} exposes {sorted(leaked)} in the schema the front renders, " + "even though it is declared as supplied by whatever runs the unit." + ) + + +def test_a_unit_can_have_no_user_configuration_at_all(units): + """Which is a real shape, not a degenerate one. + + Three units are told everything by their caller, and SaveDatasetUnit was + already configuration-free before any of this. An empty schema is valid and + renders as nothing. + """ + for name in ("SaveModelUnit", "LoadTrainingDatasetUnit", "SaveDatasetUnit"): + assert units[name]["schema"]["properties"] == {}, name + assert units[name]["configurable_object"] is True, name diff --git a/tests/back/dag/__init__.py b/tests/back/dag/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/back/dag/test_engine.py b/tests/back/dag/test_engine.py new file mode 100644 index 000000000..6975b807b --- /dev/null +++ b/tests/back/dag/test_engine.py @@ -0,0 +1,474 @@ +"""The real engine, running real units as a graph, with no database. + +The graph below is the train/test converter flow, which is the smallest thing a +single shared context genuinely cannot run: two ``LoadDatasetUnit`` nodes both +write the fixed key ``dataset``, so in one context the second load erases the +first. + + load_train --dataset--+-------------------> tx_train --dataset--> save_train + | ^ ^ + +-> fit --converter--+---+ dataset_path | + | --+ + load_test --dataset------------------------|-> tx_test --dataset--> save_test + +Two fan-outs and two joins. The number that proves it worked is 2.0: +MinMaxScaler fitted on train [0, 5, 10] learns min=0 max=10, so a test value of +20 scales to 2.0. A refit on the test data would have produced 0.0. + +The spike that established this runs the same graph against a throwaway engine +and stays as a frozen record of the unit contract. This exercises the shipped +one, which differs in the two ways that matter: nodes name their unit instead +of carrying an instance, and there are no seeds. +""" + +from pathlib import Path + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.converters.scikit_learn.min_max_scaler import MinMaxScaler +from DashAI.back.dag.engine import NullSink, artifacts_of, run +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, connect, sinks +from DashAI.back.dag.validate import instantiate, resolve_unit_class, validate +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.types.value_types import Float +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_to_path_unit import SaveDatasetToPathUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.transform_dataset_unit import TransformDatasetUnit + +FULL_SCOPE = {"columns": [], "rows": []} + + +class _Row: + """Stand-in for a Dataset ORM row, as in tests/back/units.""" + + def __init__(self, file_path): + self.file_path = file_path + self.dataset_id = None + + +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: + """A class, not a lambda: kink resolves a registered lambda by calling it.""" + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +def _write(root, **columns): + frame = pd.DataFrame(columns) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + save_dataset(to_dashai_dataset(frame, types=types), str(root / "dataset")) + return root + + +def _values(path, column="a"): + return list(load_dataset(str(path)).to_pandas()[column]) + + +@pytest.fixture(name="stored_datasets") +def fixture_stored_datasets(tmp_path): + """Dataset 7 is the training data, dataset 8 the test data. + + The engine resolves a node's unit by name, so the registry has to hold the + unit classes as well as the converter. + """ + train = _write(tmp_path / "train", a=[0.0, 5.0, 10.0]) + test = _write(tmp_path / "test", a=[20.0]) + + di["session_factory"] = _FakeSessionFactory( + {"Dataset": {7: _Row(str(train)), 8: _Row(str(test))}} + ) + di["component_registry"] = { + "MinMaxScaler": {"class": MinMaxScaler}, + "LoadDatasetUnit": {"class": LoadDatasetUnit}, + "FitConverterUnit": {"class": FitConverterUnit}, + "TransformDatasetUnit": {"class": TransformDatasetUnit}, + "SaveDatasetUnit": {"class": SaveDatasetUnit}, + "SaveDatasetToPathUnit": {"class": SaveDatasetToPathUnit}, + "SaveModelUnit": {"class": SaveModelUnit}, + "NotAUnit": {"class": MinMaxScaler}, + } + yield train, test + del di["session_factory"] + del di["component_registry"] + + +def _graph(tmp_path, converter="MinMaxScaler"): + """The diamond, wired with the bundling rule wherever it applies.""" + load_train = Node("load_train", "LoadDatasetUnit", {"dataset_id": 7}) + load_test = Node("load_test", "LoadDatasetUnit", {"dataset_id": 8}) + fit = Node( + "fit", + "FitConverterUnit", + { + "converter": {"component": converter, "params": {}}, + "scope": FULL_SCOPE, + "target": None, + }, + ) + tx_train = Node( + "tx_train", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None} + ) + tx_test = Node( + "tx_test", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None} + ) + # Saves back over the training data, so it needs the ref `dataset_path` + # from the load as well as the live dataset from the transform: a join + # whose two inputs live in different halves of the context. + save_train = Node("save_train", "SaveDatasetUnit") + save_test = Node( + "save_test", + "SaveDatasetToPathUnit", + {"path": str(tmp_path / "out" / "dataset")}, + ) + + nodes = [load_train, load_test, fit, tx_train, tx_test, save_train, save_test] + edges = [ + *connect(load_train, fit), + Edge("load_train", "dataset", "tx_train", "dataset"), + Edge("load_train", "dataset_path", "save_train", "dataset_path"), + *connect(fit, tx_train), + *connect(fit, tx_test), + Edge("load_test", "dataset", "tx_test", "dataset"), + Edge("tx_train", "dataset", "save_train", "dataset"), + *connect(tx_test, save_test), + ] + return Graph(nodes, edges) + + +def test_the_diamond_runs_and_the_fitted_state_survives_the_branch( + stored_datasets, tmp_path +): + """One fit, two branches, no refit. + + 2.0 is only reachable if the converter fitted on the train branch reached + the test branch still holding the range it learned, after crossing two + context boundaries. + """ + train, _ = stored_datasets + + run(_graph(tmp_path)) + + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + assert _values(tmp_path / "out" / "dataset") == [2.0] + + +def test_two_loads_of_the_same_unit_do_not_fight_over_the_key( + stored_datasets, tmp_path +): + """Both loads write the fixed key ``dataset`` and neither is lost. + + This is what a shared context cannot do, and the reason each node gets one + of its own. + """ + train, _ = stored_datasets + + run(_graph(tmp_path)) + + # A second load overwriting the first would have scaled the test value + # against the test data's own range, giving 0.0. + assert _values(tmp_path / "out" / "dataset") == [2.0] + assert _values(train / "dataset") == [0.0, 0.5, 1.0] + + +def test_a_node_names_its_unit_and_the_engine_builds_one_instance_each( + stored_datasets, tmp_path +): + """Each node owns its instance, which is what per-instance state needs. + + A unit memoizes work on itself -- ``FitConverterUnit`` its converter + class -- so two nodes sharing an instance would have the second silently + run with the first one's. + """ + a = instantiate( + Node("a", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + ) + b = instantiate( + Node("b", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + ) + + assert a is not b + assert type(a) is type(b) is TransformDatasetUnit + + +def test_the_execution_order_is_topological_not_the_order_given( + stored_datasets, tmp_path +): + """Order comes from the edges, not from the sequence the nodes arrived in. + + Its predecessor iterated the steps in whatever order the front sent them -- + creation order on the canvas -- and never read the edges it persisted. + """ + graph = _graph(tmp_path) + reversed_graph = Graph(list(reversed(graph.nodes)), graph.edges) + + order = validate(reversed_graph) + + assert order.index("load_train") < order.index("fit") + assert order.index("fit") < order.index("tx_test") + assert order.index("tx_train") < order.index("save_train") + + +def test_a_missing_input_is_caught_before_anything_runs(stored_datasets): + save = Node("save", "SaveDatasetUnit") + + with pytest.raises(GraphError, match="requires 'dataset'"): + validate(Graph([save], [])) + + +def test_two_edges_into_one_port_are_rejected(stored_datasets, tmp_path): + """The failure its predecessor had: a merge where the last writer wins.""" + graph = _graph(tmp_path) + doubled = Graph( + graph.nodes, + [*graph.edges, Edge("load_test", "dataset", "tx_train", "dataset")], + ) + + with pytest.raises(GraphError, match="from 2 edges"): + validate(doubled) + + +def test_a_cycle_is_reported(stored_datasets): + a = Node("a", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + b = Node("b", "TransformDatasetUnit", {"scope": FULL_SCOPE, "target": None}) + cyclic = Graph( + [a, b], + [ + Edge("a", "dataset", "b", "dataset"), + Edge("b", "dataset", "a", "dataset"), + ], + ) + + with pytest.raises(GraphError, match="cycle"): + validate(cyclic) + + +def test_an_edge_naming_a_key_the_unit_does_not_declare_is_rejected( + stored_datasets, tmp_path +): + graph = _graph(tmp_path) + bogus = Graph( + graph.nodes, + [*graph.edges, Edge("load_train", "nonsense", "save_train", "dataset_path")], + ) + + with pytest.raises(GraphError, match="does not provide 'nonsense'"): + validate(bogus) + + +def test_an_unknown_unit_name_is_rejected(stored_datasets): + with pytest.raises(GraphError, match="no unit named 'Nope'"): + validate(Graph([Node("x", "Nope")], [])) + + +def test_a_registered_component_that_is_not_a_unit_is_rejected(stored_datasets): + """The registry holds every kind of component, not only units.""" + with pytest.raises(GraphError, match="not a unit"): + validate(Graph([Node("x", "NotAUnit")], [])) + + +def test_duplicate_node_ids_are_rejected(stored_datasets, tmp_path): + twice = Node("same", "SaveDatasetToPathUnit", {"path": str(tmp_path / "d")}) + with pytest.raises(GraphError, match="Duplicate node ids"): + validate(Graph([twice, twice], [])) + + +def test_every_problem_is_reported_at_once(stored_datasets): + """A user fixing a graph wants the whole list, not one error per attempt.""" + graph = Graph( + [Node("save", "SaveDatasetUnit"), Node("save", "SaveDatasetUnit")], [] + ) + + with pytest.raises(GraphError) as caught: + validate(graph) + + message = str(caught.value) + assert "Duplicate node ids" in message + assert "requires 'dataset'" in message + + +def test_moving_a_value_across_an_edge_keeps_the_half_it_came_from( + stored_datasets, tmp_path +): + """Refs travel as refs, live objects as live objects. + + ``dataset_path`` is a reference and ``dataset`` is a cached object, and the + two have incompatible rules: a dataset handed to ``put_ref`` raises, and a + reference handed to ``put`` drops the copy-on-read guarantee it depended + on. + """ + contexts = run(_graph(tmp_path)) + + save_train = contexts["save_train"] + assert save_train.origin("dataset_path") == "ref" + assert save_train.origin("dataset") == "cache" + + +def test_a_context_is_released_once_nothing_downstream_reads_it( + stored_datasets, tmp_path +): + """Intermediate datasets do not stay alive for the length of the run.""" + graph = _graph(tmp_path) + + contexts = run(graph) + + assert set(contexts) == sinks(graph) + assert "load_train" not in contexts + + +def test_artifacts_are_the_serializable_outputs_only(stored_datasets, tmp_path): + """What gets recorded is the reference half of PROVIDES. + + The cache holds live datasets, models and tasks: not serializable, and + always derivable again. So ``LoadDatasetUnit`` records the id and the path + it published rather than the dataset itself, and ``TransformDatasetUnit``, + whose only output is the live dataset, records nothing -- a real case, not + a gap. + """ + train, _ = stored_datasets + sink = _RecordingSink() + + run(_graph(tmp_path), sink) + + recorded = {event[1]: event[2] for event in sink.events if event[0] == "finished"} + + assert set(recorded["load_train"]) == {"dataset_id", "dataset_path"} + assert recorded["load_train"]["dataset_id"] == 7 + # Compared as a path, not as text: the unit joins with a forward slash and + # Path renders a backslash on Windows. + assert Path(recorded["load_train"]["dataset_path"]) == train / "dataset" + assert recorded["tx_train"] == {} + + +def test_artifacts_of_reads_only_the_reference_half(stored_datasets): + """The same rule, stated directly against a hand-built context.""" + ctx = ExecutionContext() + ctx.put_ref("dataset_path", "/tmp/ds") + ctx.put_ref("dataset_id", 7) + ctx.put("dataset", object()) + + assert artifacts_of(resolve_unit_class("LoadDatasetUnit"), ctx) == { + "dataset_path": "/tmp/ds", + "dataset_id": 7, + } + + +def test_the_bundling_rule_picks_the_keys_two_nodes_agree_on(stored_datasets): + """One drawn edge between two nodes stands for a set of keys.""" + fit = Node("fit", "FitConverterUnit", {}) + tx = Node("tx", "TransformDatasetUnit", {}) + + assert connect(fit, tx) == ( + Edge("fit", "fitted_converter", "tx", "fitted_converter"), + ) + + +def test_a_node_belongs_to_itself_when_it_has_no_block(): + """Every node has a block, so the tracking never has a null to handle.""" + alone = Node("solo", "SaveDatasetUnit") + grouped = Node("inner", "SaveDatasetUnit", {}, block_id="train-1") + + assert alone.block_id == "solo" + assert grouped.block_id == "train-1" + + +class _RecordingSink(NullSink): + """Records the notifications the engine sends, in order.""" + + def __init__(self): + self.events = [] + + def run_started(self, order): + self.events.append(("run_started", tuple(order))) + + def node_started(self, node_id, payload): + self.events.append(("started", node_id, dict(payload))) + + def node_finished(self, node_id, artifacts, payload): + self.events.append(("finished", node_id, dict(artifacts))) + + def node_failed(self, node_id, message): + self.events.append(("failed", node_id)) + + def nodes_cancelled(self, node_ids): + self.events.append(("cancelled", tuple(node_ids))) + + def run_finished(self): + self.events.append(("run_finished",)) + + def run_failed(self, message): + self.events.append(("run_failed",)) + + +def test_the_sink_sees_every_node_start_and_finish(stored_datasets, tmp_path): + sink = _RecordingSink() + + run(_graph(tmp_path), sink) + + kinds = [event[0] for event in sink.events] + assert kinds[0] == "run_started" + assert kinds[-1] == "run_finished" + assert kinds.count("started") == 7 + assert kinds.count("finished") == 7 + assert "failed" not in kinds + + +def test_a_failing_node_cancels_everything_after_it(stored_datasets, tmp_path): + """A run that died halfway must not look like one still in flight.""" + graph = _graph(tmp_path, converter="NoSuchConverter") + sink = _RecordingSink() + + # The unit reports an unresolvable converter as a JobError, and the engine + # lets it through rather than wrapping it: the message the user reads is + # the unit's own. + with pytest.raises(JobError, match="NoSuchConverter"): + run(graph, sink) + + failed = [event for event in sink.events if event[0] == "failed"] + cancelled = [event for event in sink.events if event[0] == "cancelled"] + + assert [event[1] for event in failed] == ["fit"] + assert cancelled, "the nodes after the failure have to be reported" + assert "tx_train" in cancelled[0][1] + assert "save_test" in cancelled[0][1] + assert ("run_failed",) in sink.events + assert ("run_finished",) not in sink.events + + +def test_nothing_runs_when_the_graph_does_not_validate(stored_datasets, tmp_path): + """Validation comes first, so a broken graph costs no work.""" + sink = _RecordingSink() + save = Node("save", "SaveDatasetUnit") + + with pytest.raises(GraphError): + run(Graph([save], []), sink) + + assert sink.events == [] diff --git a/tests/back/dag/test_expand.py b/tests/back/dag/test_expand.py new file mode 100644 index 000000000..92bfbbc1e --- /dev/null +++ b/tests/back/dag/test_expand.py @@ -0,0 +1,285 @@ +"""Expanding what a canvas holds into the graph the engine runs.""" + +import pytest +from kink import di + +from DashAI.back.dag.expand import artifact_prefix, expand +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node +from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_model_unit import SaveModelUnit + + +@pytest.fixture(name="registry", autouse=True) +def fixture_registry(): + di["component_registry"] = { + "LoadDatasetUnit": {"class": LoadDatasetUnit}, + "SaveDatasetUnit": {"class": SaveDatasetUnit}, + "SaveModelUnit": {"class": SaveModelUnit}, + } + yield + del di["component_registry"] + + +def _block(block_id, unit, config=None, unit_id=None): + return { + "id": block_id, + "type": f"{unit}Block", + "units": [{"id": unit_id or block_id, "unit": unit, "config": config or {}}], + } + + +def test_one_unit_per_block_expands_to_the_graph_written_by_hand(): + steps = [ + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + _block("save-1", "SaveDatasetUnit"), + ] + edges = [{"source": "load-1", "target": "save-1"}] + + graph = expand(steps, edges, pipeline_run_id=3) + + assert graph == Graph( + [ + Node("load-1", "LoadDatasetUnit", {"dataset_id": 7}, block_id="load-1"), + Node("save-1", "SaveDatasetUnit", {}, block_id="save-1"), + ], + [ + Edge("load-1", "dataset", "save-1", "dataset"), + Edge("load-1", "dataset_path", "save-1", "dataset_path"), + ], + ) + + +def test_one_drawn_edge_carries_every_key_the_two_units_agree_on(): + """A canvas cannot draw one edge per key, so one stands for the set.""" + steps = [ + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + _block("save-1", "SaveDatasetUnit"), + ] + + graph = expand(steps, [{"source": "load-1", "target": "save-1"}], 3) + + assert {(edge.src_key, edge.dst_key) for edge in graph.edges} == { + ("dataset", "dataset"), + ("dataset_path", "dataset_path"), + } + + +def test_a_node_keeps_the_block_it_came_from(): + """The canvas colours a block from the status of its nodes.""" + steps = [_block("train-1", "SaveModelUnit", unit_id="train-1/save")] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].id == "train-1/save" + assert graph.nodes[0].block_id == "train-1" + + +def test_the_engine_names_the_artifacts_of_a_unit_that_takes_a_prefix(): + steps = [_block("save-1", "SaveModelUnit")] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].config["artifact_prefix"] == "pipeline-3-save-1" + + +def test_a_prefix_stored_in_the_graph_is_discarded(): + """No stored value can be right, so none is honoured. + + The prefix is built from the id of the run, and there is no run when a node + is configured -- so anything already there was chosen without the one thing + that decides it. Overriding is also what makes two nodes sharing a prefix + impossible rather than merely unlikely. + """ + steps = [_block("save-1", "SaveModelUnit", {"artifact_prefix": "mine"})] + + graph = expand(steps, [], pipeline_run_id=3) + + assert graph.nodes[0].config["artifact_prefix"] == "pipeline-3-save-1" + + +def test_two_saving_nodes_cannot_end_up_with_the_same_prefix(): + """Which is why the collision needs no validator to catch it.""" + steps = [ + _block("save-a", "SaveModelUnit", {"artifact_prefix": "same"}), + _block("save-b", "SaveModelUnit", {"artifact_prefix": "same"}), + ] + + prefixes = {node.config["artifact_prefix"] for node in expand(steps, [], 3).nodes} + + assert len(prefixes) == 2 + + +@pytest.mark.parametrize( + ("first", "second"), + [("save.a", "save_a"), ("train-1/build", "train-1_build"), ("a b", "a_b")], +) +def test_sanitizing_two_different_ids_cannot_produce_one_prefix(first, second): + """The escape has to be reversible, not merely safe. + + Replacing every unsafe character with a plain ``_`` is safe but collapses + distinct ids: ``save.a`` and ``save_a`` would land on the same directory and + the second saving node would overwrite the first one's model in silence. + That is why ``_`` is escaped along with everything else. + """ + assert artifact_prefix(3, first) != artifact_prefix(3, second) + + +def test_a_unit_that_takes_no_prefix_does_not_get_one(): + steps = [_block("load-1", "LoadDatasetUnit", {"dataset_id": 7})] + + graph = expand(steps, [], pipeline_run_id=3) + + assert "artifact_prefix" not in graph.nodes[0].config + + +def test_two_runs_of_one_pipeline_name_their_artifacts_apart(): + """Otherwise the second run would write over the first one's model.""" + steps = [_block("save-1", "SaveModelUnit")] + + first = expand(steps, [], pipeline_run_id=3).nodes[0].config["artifact_prefix"] + second = expand(steps, [], pipeline_run_id=4).nodes[0].config["artifact_prefix"] + + assert first != second + + +def test_the_prefix_can_never_be_read_as_a_path(): + """A node id a user chose has to be brought into the safe alphabet. + + ``os.path.join`` with a separator in the prefix produces a destination + outside the runs directory, and ``SaveModelUnit.validate`` refuses one -- + so an id with a slash in it has to be sanitized here rather than failing at + the node. + """ + # Each unsafe character becomes _, so the mapping is reversible. + assert artifact_prefix(3, "train-1/save") == "pipeline-3-train-1_2fsave" + # The dot is escaped too, which is what leaves a parent reference inert + # rather than merely separator-free. + assert artifact_prefix(3, "../escape") == "pipeline-3-_2e_2e_2fescape" + + steps = [_block("save-1", "SaveModelUnit", unit_id="train-1/save")] + prefix = expand(steps, [], 3).nodes[0].config["artifact_prefix"] + + # And what comes out passes the guard at the unit. + SaveModelUnit(artifact_prefix=prefix).validate(None) + + +def test_a_pipeline_run_id_alone_would_collide_with_a_real_run(): + """The whole reason the prefix is not just the id. + + RUNS_PATH is shared with every real Run, and pipeline_run.id and run.id are + independent sequences that both start at 1. + """ + assert artifact_prefix(3, "save-1") != "3" + assert artifact_prefix(3, "save-1").startswith("pipeline-") + + +def test_a_block_with_no_units_is_refused_by_name(): + """Rows written by the previous subsystem land here. + + Their steps name a node type and carry a single config, with no unit to + resolve, so they are refused with a message that says so rather than + failing somewhere deeper. + """ + old_style = [{"id": "DataSelector-1", "type": "DataSelector", "config": {}}] + + with pytest.raises(GraphError, match="previous pipeline subsystem"): + expand(old_style, [], 3) + + +def test_more_than_one_unit_per_block_is_refused_rather_than_guessed(): + """Wiring across a sequence boundary has more than one defensible answer.""" + steps = [ + { + "id": "train-1", + "units": [ + {"id": "a", "unit": "LoadDatasetUnit", "config": {"dataset_id": 7}}, + {"id": "b", "unit": "SaveDatasetUnit", "config": {}}, + ], + } + ] + + with pytest.raises(GraphError, match="one unit per block"): + expand(steps, [], 3) + + +def test_a_block_without_an_id_is_refused(): + with pytest.raises(GraphError, match="no id"): + expand([{"units": []}], [], 3) + + +def test_two_blocks_with_the_same_id_are_refused(): + steps = [ + _block("same", "LoadDatasetUnit", {"dataset_id": 7}), + _block("same", "SaveDatasetUnit"), + ] + + with pytest.raises(GraphError, match="more than one block"): + expand(steps, [], 3) + + +def test_an_edge_to_a_block_that_is_not_there_is_refused(): + steps = [_block("load-1", "LoadDatasetUnit", {"dataset_id": 7})] + + with pytest.raises(GraphError, match="not a block"): + expand(steps, [{"source": "load-1", "target": "ghost"}], 3) + + +def test_an_edge_that_would_carry_nothing_is_refused(): + """Two units with no key in common cannot be usefully connected.""" + steps = [ + _block("save-1", "SaveModelUnit"), + _block("load-1", "LoadDatasetUnit", {"dataset_id": 7}), + ] + + with pytest.raises(GraphError, match="would carry nothing"): + expand(steps, [{"source": "save-1", "target": "load-1"}], 3) + + +def test_a_node_that_takes_a_run_id_is_told_there_is_no_run(): + """A pipeline has no Run row, and that is a value rather than an omission. + + It is how a node says the model it builds belongs to no run, which is what + keeps that model from trying to log metrics against a foreign key that + points at nothing. + """ + di["component_registry"]["FitModelUnit"] = {"class": FitModelUnit} + steps = [ + _block( + "fit-1", + "FitModelUnit", + { + "optimizer": {"component": "OptunaOptimizer", "params": {}}, + "goal_metric": "Accuracy", + }, + ) + ] + + config = expand(steps, [], pipeline_run_id=3).nodes[0].config + + assert config["run_id"] is None + assert config["artifact_prefix"] == "pipeline-3-fit-1" + + +def test_a_run_id_stored_in_the_graph_is_discarded(): + """A pipeline has no run, so a stored one is overridden rather than trusted. + + This is what makes the sandbox mechanical instead of conventional: a + hand-edited row cannot point a training node at a real run and start + writing Metric rows into it. + """ + di["component_registry"]["FitModelUnit"] = {"class": FitModelUnit} + steps = [ + _block( + "fit-1", + "FitModelUnit", + { + "optimizer": {"component": "OptunaOptimizer", "params": {}}, + "goal_metric": "Accuracy", + "run_id": 17, + }, + ) + ] + + assert expand(steps, [], 3).nodes[0].config["run_id"] is None diff --git a/tests/back/dag/test_pipeline_job.py b/tests/back/dag/test_pipeline_job.py new file mode 100644 index 000000000..8c1327c1a --- /dev/null +++ b/tests/back/dag/test_pipeline_job.py @@ -0,0 +1,339 @@ +"""PipelineJob end to end, against a real database. + +Its predecessor implemented two of BaseJob's four abstract methods, so it could +not be instantiated at all: asking the job endpoint for a PipelineJob raised +TypeError. These check the job runs, and that the four methods are there. +""" + +import pytest +from kink import di +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import StaticPool + +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dependencies.database.models import ( + Base, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, +) +from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.job.pipeline_job import PipelineJob +from DashAI.back.units.base_unit import BaseUnit + + +class _EmitUnit(BaseUnit): + REQUIRES = () + PROVIDES = ("dataset_path", "dataset") + + def execute(self, ctx): + ctx.put_ref("dataset_path", "/tmp/ds") + ctx.put("dataset", object()) + + +class _ConsumeUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = ("results_path",) + + def execute(self, ctx): + ctx.require("dataset") + ctx.put_ref("results_path", "/tmp/results.json") + + +class _FailUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx): + raise JobError("this node was always going to fail") + + +@pytest.fixture(name="database") +def fixture_database(): + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine) + + class _Factory: + def __call__(self): + return factory() + + di["session_factory"] = _Factory() + di["component_registry"] = { + "EmitUnit": {"class": _EmitUnit}, + "ConsumeUnit": {"class": _ConsumeUnit}, + "FailUnit": {"class": _FailUnit}, + } + yield factory + del di["session_factory"] + del di["component_registry"] + engine.dispose() + + +def _pipeline(factory, steps, edges, name="a pipeline"): + with factory() as db: + pipeline = Pipeline(name=name, steps=steps, edges=edges) + db.add(pipeline) + db.commit() + return pipeline.id + + +def _block(block_id, unit): + return { + "id": block_id, + "units": [{"id": block_id, "unit": unit, "config": {}}], + } + + +def test_the_job_implements_every_abstract_method_of_base_job(): + """The reason its predecessor could not be instantiated at all.""" + missing = { + name + for name in getattr(BaseJob, "__abstractmethods__", set()) + if getattr(PipelineJob, name) is getattr(BaseJob, name) + } + assert not missing + assert PipelineJob(pipeline_id=1) is not None + + +def test_a_pipeline_runs_and_records_what_each_node_produced(database): + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("consume-1", "ConsumeUnit")], + [{"source": "emit-1", "target": "consume-1"}], + ) + + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.FINISHED + assert {row.node_id for row in pipeline_run.node_runs} == { + "emit-1", + "consume-1", + } + assert {row.status for row in pipeline_run.node_runs} == { + NodeRunStatus.FINISHED + } + artifacts = { + (row.node_run.node_id, row.key) for row in db.query(NodeArtifact).all() + } + + assert artifacts == { + ("emit-1", "dataset_path"), + ("consume-1", "results_path"), + } + + +def test_the_run_freezes_the_expanded_graph(database): + """A run records units and keys, not the blocks a canvas drew.""" + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("consume-1", "ConsumeUnit")], + [{"source": "emit-1", "target": "consume-1"}], + ) + + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + + assert pipeline_run.edges == [ + { + "src": "emit-1", + "src_key": "dataset", + "dst": "consume-1", + "dst_key": "dataset", + } + ] + assert [step["unit"] for step in pipeline_run.steps] == [ + "EmitUnit", + "ConsumeUnit", + ] + + +def test_a_graph_that_cannot_work_says_why_and_marks_the_run(database): + """An unfed input is a mistake in the graph, reported as one.""" + pipeline_id = _pipeline(database, [_block("consume-1", "ConsumeUnit")], []) + + with pytest.raises(JobError, match="requires 'dataset'"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "requires 'dataset'" in pipeline_run.error_message + # Nothing ran, so there is nothing to show per node. + assert db.query(NodeRun).count() == 0 + + +def test_steps_from_the_previous_subsystem_are_refused_by_name(database): + """Rows the old canvas saved name a node type and have no unit.""" + pipeline_id = _pipeline( + database, + [{"id": "DataSelector-1", "type": "DataSelector", "config": {}}], + [], + ) + + with pytest.raises(JobError, match="previous pipeline subsystem"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.ERROR + + +def test_a_pipeline_with_no_steps_is_refused_before_a_run_is_created(database): + pipeline_id = _pipeline(database, [], []) + + with pytest.raises(JobError, match="no steps"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).count() == 0 + + +def test_a_missing_pipeline_is_refused(database): + with pytest.raises(JobError, match="does not exist"): + PipelineJob(pipeline_id=9999).run() + + +def test_two_runs_of_one_pipeline_name_their_artifacts_apart(database): + """The reason the run row is created before the graph is expanded.""" + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + PipelineJob(pipeline_id=pipeline_id).run() + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + runs = db.query(PipelineRun).order_by(PipelineRun.id).all() + assert len(runs) == 2 + assert runs[0].id != runs[1].id + # Both runs recorded their own history rather than overwriting one. + assert {row.status for run in runs for row in run.node_runs} == { + NodeRunStatus.FINISHED + } + + +def test_the_job_name_comes_from_the_pipeline(database): + pipeline_id = _pipeline(database, [], [], name="my flow") + + assert PipelineJob(pipeline_id=pipeline_id).get_job_name() == "Pipeline: my flow" + assert PipelineJob(pipeline_id=9999).get_job_name() == "Pipeline (9999)" + + +def test_the_job_kwargs_are_only_plain_data(database): + """The whole job is serialized with dill to reach the worker process.""" + import dill + + job = PipelineJob(pipeline_id=7) + + assert job.kwargs == {"pipeline_id": 7} + assert dill.loads(dill.dumps(job)).kwargs == {"pipeline_id": 7} + + +def test_the_id_the_front_actually_sends_is_accepted(database): + """Regression: the wire contract predates this job. + + ``api/job.ts`` posts ``kwargs: {id: pipelineId}``, and ``jobs.py`` passes + kwargs through verbatim. Reading only ``pipeline_id`` meant every enqueued + pipeline died in the worker with a raw KeyError, before any run row existed. + """ + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + PipelineJob(id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.FINISHED + + +def test_either_name_reaches_the_same_pipeline(database): + pipeline_id = _pipeline(database, [], [], name="named") + + assert PipelineJob(id=pipeline_id).get_job_name() == "Pipeline: named" + assert PipelineJob(pipeline_id=pipeline_id).get_job_name() == "Pipeline: named" + + +def test_a_job_with_no_id_at_all_says_so(database): + with pytest.raises(JobError, match="No pipeline id"): + PipelineJob().run() + + assert PipelineJob().get_job_name() == "Pipeline" + + +def test_a_run_is_never_left_looking_like_it_is_still_going(database, monkeypatch): + """The backstop for the paths where the sink recorded nothing. + + The sink's own calls sit outside the engine's try, so one of them failing + -- a locked database, a run row deleted underneath -- raises without any + status having been written. A run stuck in STARTED with no error message is + indistinguishable from one still in flight. + """ + from DashAI.back.dag import tracking + + def explode(self, node_id, payload): + raise RuntimeError("the sink could not write") + + monkeypatch.setattr(tracking.DatabaseSink, "node_started", explode) + + pipeline_id = _pipeline(database, [_block("emit-1", "EmitUnit")], []) + + with pytest.raises(JobError, match="could not write"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "could not write" in pipeline_run.error_message + + +def test_the_backstop_does_not_overwrite_what_the_sink_recorded(database): + """A node's own message is the better one, so it has to survive.""" + pipeline_id = _pipeline( + database, + [_block("emit-1", "EmitUnit"), _block("boom-1", "FailUnit")], + [{"source": "emit-1", "target": "boom-1"}], + ) + + with pytest.raises(JobError, match="always going to fail"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + pipeline_run = db.query(PipelineRun).one() + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "always going to fail" in pipeline_run.error_message + + +class _NeedsARuntimeParam(BaseUnit): + """Stands in for a unit only a particular job knows how to configure.""" + + REQUIRES = () + PROVIDES = ("dataset_path",) + RUNTIME_PARAMS = ("temp_path",) + + def execute(self, ctx): + ctx.put_ref("dataset_path", self.config["temp_path"]) + + +def test_a_runtime_param_nobody_supplies_fails_before_anything_runs(database): + """Not halfway through, after earlier nodes already wrote to disk. + + Only two runtime params are about a pipeline run, so the engine can answer + those. A unit whose runtime params only a particular job knows how to fill + is not usable as a node yet, and this is where a user is told so. + """ + di["component_registry"]["NeedsARuntimeParam"] = {"class": _NeedsARuntimeParam} + pipeline_id = _pipeline(database, [_block("odd-1", "NeedsARuntimeParam")], []) + + with pytest.raises(JobError, match=r"needs \['temp_path'\]"): + PipelineJob(pipeline_id=pipeline_id).run() + + with database() as db: + assert db.query(PipelineRun).one().status == PipelineRunStatus.ERROR + # Nothing ran, so there is nothing to show per node. + assert db.query(NodeRun).count() == 0 diff --git a/tests/back/dag/test_tracking.py b/tests/back/dag/test_tracking.py new file mode 100644 index 000000000..ec1c6b885 --- /dev/null +++ b/tests/back/dag/test_tracking.py @@ -0,0 +1,266 @@ +"""What a run writes down, against a real database. + +The engine is exercised with a sink that records nothing in test_engine.py; +this checks the sink that persists, including the two things "no row" cannot +express: a node still waiting its turn, and a node whose run died before +reaching it. +""" + +import pytest +from kink import di +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker +from sqlalchemy.pool import StaticPool + +from DashAI.back.core.enums.status import NodeRunStatus, PipelineRunStatus +from DashAI.back.dag.engine import run +from DashAI.back.dag.graph import Edge, Graph, GraphError, Node, dump, load +from DashAI.back.dag.tracking import DatabaseSink +from DashAI.back.dependencies.database.models import ( + Base, + NodeArtifact, + NodeRun, + Pipeline, + PipelineRun, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit + + +class _EmitUnit(BaseUnit): + """Publishes a reference and a live object, so both halves are covered.""" + + REQUIRES = () + PROVIDES = ("dataset_path", "dataset") + + def execute(self, ctx): + ctx.put_ref("dataset_path", self.config["path"]) + ctx.put("dataset", object()) + + +class _ConsumeUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = ("results_path",) + + def execute(self, ctx): + ctx.require("dataset") + ctx.put_ref("results_path", "/tmp/results.json") + + +class _FailUnit(BaseUnit): + REQUIRES = ("dataset",) + PROVIDES = () + + def execute(self, ctx): + raise JobError("this node was always going to fail") + + +@pytest.fixture(name="database") +def fixture_database(): + """An in-memory database shared across sessions. + + StaticPool because every notification opens its own session: without it + each connection would get a fresh empty database. + """ + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + Base.metadata.create_all(engine) + factory = sessionmaker(bind=engine) + + class _Factory: + """Not a lambda: kink calls a registered lambda when resolving it.""" + + def __call__(self): + return factory() + + di["session_factory"] = _Factory() + di["component_registry"] = { + "EmitUnit": {"class": _EmitUnit}, + "ConsumeUnit": {"class": _ConsumeUnit}, + "FailUnit": {"class": _FailUnit}, + } + + with factory() as db: + pipeline = Pipeline(name="a pipeline", steps=[], edges=[]) + db.add(pipeline) + db.commit() + # The run row is the caller's to create: naming a node's artifacts + # needs its id, and that comes before the graph exists. + pipeline_run = PipelineRun(pipeline_id=pipeline.id) + db.add(pipeline_run) + db.commit() + pipeline_run_id = pipeline_run.id + + yield factory, pipeline_run_id + + del di["session_factory"] + del di["component_registry"] + engine.dispose() + + +def _pair_graph(): + emit = Node("emit", "EmitUnit", {"path": "/tmp/ds"}) + consume = Node("consume", "ConsumeUnit") + return Graph([emit, consume], [Edge("emit", "dataset", "consume", "dataset")]) + + +def test_a_finished_run_records_every_node_and_its_artifacts(database): + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.FINISHED + assert pipeline_run.start_time is not None + assert pipeline_run.end_time is not None + + node_runs = {row.node_id: row for row in pipeline_run.node_runs} + assert set(node_runs) == {"emit", "consume"} + for row in node_runs.values(): + assert row.status == NodeRunStatus.FINISHED + assert row.start_time is not None + assert row.end_time is not None + + artifacts = { + (row.node_run.node_id, row.key): row.value + for row in db.query(NodeArtifact).all() + } + + # The live dataset is not recorded: the cache half is not serializable and + # is always derivable again. The path it published is. + assert artifacts == { + ("emit", "dataset_path"): "/tmp/ds", + ("consume", "results_path"): "/tmp/results.json", + } + + +def test_the_run_freezes_the_graph_it_executed(database): + """A past run stays readable after the pipeline it came from is edited.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + frozen_steps, frozen_edges = pipeline_run.steps, pipeline_run.edges + + steps, edges = dump(graph) + assert frozen_steps == steps + assert frozen_edges == edges + # And the frozen form is enough to rebuild the graph, without re-running + # whatever produced it. + assert load(frozen_steps, frozen_edges) == graph + + +def test_a_failure_marks_the_node_and_cancels_the_rest(database): + """CANCELLED is what NOT_STARTED cannot say. + + Without it, a run that died before reaching a node is indistinguishable + from one still in flight. + """ + factory, pipeline_run_id = database + emit = Node("emit", "EmitUnit", {"path": "/tmp/ds"}) + boom = Node("boom", "FailUnit") + consume = Node("consume", "ConsumeUnit") + graph = Graph( + [emit, boom, consume], + [ + Edge("emit", "dataset", "boom", "dataset"), + Edge("emit", "dataset", "consume", "dataset"), + ], + ) + sink = DatabaseSink(pipeline_run_id, graph) + + with pytest.raises(JobError, match="always going to fail"): + run(graph, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, sink.pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.ERROR + assert "always going to fail" in pipeline_run.error_message + + rows = {row.node_id: row for row in pipeline_run.node_runs} + + assert rows["emit"].status == NodeRunStatus.FINISHED + assert rows["boom"].status == NodeRunStatus.ERROR + assert "always going to fail" in rows["boom"].error_message + assert rows["consume"].status == NodeRunStatus.CANCELLED + + +def test_every_node_has_a_row_before_the_first_one_runs(database): + """Rows exist up front, so a node waiting its turn is visible as waiting.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + sink.run_started(["emit", "consume"]) + + with factory() as db: + rows = db.query(NodeRun).all() + assert {row.node_id for row in rows} == {"emit", "consume"} + assert {row.status for row in rows} == {NodeRunStatus.NOT_STARTED} + + +def test_a_node_belongs_to_a_block_even_when_it_is_alone(database): + """block_id is never null, so grouping later needs no migration.""" + factory, pipeline_run_id = database + graph = Graph( + [Node("inner", "EmitUnit", {"path": "/tmp/ds"}, block_id="train-1")], [] + ) + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + row = db.query(NodeRun).one() + assert row.node_id == "inner" + assert row.block_id == "train-1" + assert row.node_type == "EmitUnit" + + +def test_a_graph_that_does_not_validate_leaves_the_run_untouched(database): + """Validation happens before the engine reports anything. + + So the run stays exactly as its caller made it -- no node rows, no start + time, still NOT_STARTED -- and saying why is the caller's to do. Marking it + started and then failed would claim a run began when nothing did. + """ + factory, pipeline_run_id = database + orphan = Graph([Node("consume", "ConsumeUnit")], []) + sink = DatabaseSink(pipeline_run_id, orphan) + + with pytest.raises(GraphError, match="requires 'dataset'"): + run(orphan, sink) + + with factory() as db: + pipeline_run = db.get(PipelineRun, pipeline_run_id) + assert pipeline_run.status == PipelineRunStatus.NOT_STARTED + assert pipeline_run.start_time is None + assert db.query(NodeRun).count() == 0 + + +def test_the_node_input_and_output_are_the_serializable_halves(database): + """What crossed into a node and what it left behind, as plain data.""" + factory, pipeline_run_id = database + graph = _pair_graph() + sink = DatabaseSink(pipeline_run_id, graph) + + run(graph, sink) + + with factory() as db: + rows = {row.node_id: row for row in db.query(NodeRun).all()} + + # `dataset` is a live object, so it appears in neither: this is what + # replaced skipping the key by name when serializing. + assert rows["emit"].input == {} + assert rows["emit"].output == {"dataset_path": "/tmp/ds"} + assert rows["consume"].input == {} + assert rows["consume"].output == {"results_path": "/tmp/results.json"} diff --git a/tests/back/pipeline_spike/test_dag_engine_spike.py b/tests/back/pipeline_spike/test_dag_engine_spike.py index c74079077..797a88968 100644 --- a/tests/back/pipeline_spike/test_dag_engine_spike.py +++ b/tests/back/pipeline_spike/test_dag_engine_spike.py @@ -234,25 +234,33 @@ def test_a_cycle_is_reported(): validate(graph) -def test_run_id_has_to_be_seeded_because_no_unit_publishes_it(): - """``run_id`` is an orphan input: four units require it, none provides it. - - In a job it arrived through ``self.kwargs``. A graph has to inject it from - outside, and the validator cannot tell that injection apart from a wire the - user forgot to draw. +def test_run_id_is_no_longer_an_orphan_input(): + """``run_id`` used to be a context key no unit published. + + Four units required it and none provided it, so a graph had to inject it + from outside and the validator could not tell that injection apart from a + wire the user forgot to draw. It was never data flowing through the graph: + in a job it arrived through ``self.kwargs``. + + It is configuration now, which is what closed that hole. Every key left in + ``REQUIRES`` has a unit that publishes it, so "nothing supplies this" means + a missing edge and nothing else. The seeds mechanism this spike used for it + has no remaining caller. """ - save = Node("save", SaveModelUnit()) + save = Node("save", SaveModelUnit(artifact_prefix="pipeline-1-save")) - with pytest.raises(GraphError, match="save requires 'run_id'"): + with pytest.raises(GraphError, match="save requires 'model'"): validate(Graph([save], [])) - seeded = Node("save", SaveModelUnit(), seeds={"run_id": 1}) - with pytest.raises(GraphError, match="save requires 'model'"): - validate(Graph([seeded], [])) + assert "run_id" not in SaveModelUnit.REQUIRES def test_a_seed_for_a_key_the_unit_never_uses_is_rejected(): - save = Node("save", SaveModelUnit(), seeds={"run_id": 1, "nonsense": 2}) + save = Node( + "save", + SaveModelUnit(artifact_prefix="pipeline-1-save"), + seeds={"nonsense": 2}, + ) with pytest.raises(GraphError, match="seeded with 'nonsense'"): validate(Graph([save], [])) diff --git a/tests/back/units/test_build_model_unit.py b/tests/back/units/test_build_model_unit.py index a14aee9ab..6e17fbb4a 100644 --- a/tests/back/units/test_build_model_unit.py +++ b/tests/back/units/test_build_model_unit.py @@ -29,12 +29,13 @@ def fake_registry(): del di["component_registry"] -def _build_unit(model_name): +def _build_unit(model_name, run_id=1): return BuildModelUnit( model={"component": model_name, "params": {}}, train_metrics=[], validation_metrics=[], test_metrics=[], + run_id=run_id, ) @@ -51,7 +52,6 @@ def test_two_build_model_units_in_one_context_resolve_independently(fake_registr ctx.put("x", {"train": None, "validation": None}) ctx.put("y", {"train": None, "validation": None}) ctx.put("n_labels", None) - ctx.put_ref("run_id", 1) ctx.put_ref("task_name", "ATask") a = _build_unit("ModelA") @@ -74,3 +74,39 @@ def test_resolve_model_class_is_memoized_per_instance_not_shared(fake_registry): assert b._resolve_model_class() is _ModelB # Calling again must return the same, still-correct class from the cache. assert a._resolve_model_class() is _ModelA + + +def test_a_run_id_of_none_leaves_the_model_detached_from_any_run(fake_registry): + """No run id means no run, and a model with no run logs no metrics. + + ``ModelFactory`` hangs the run id on the model instance, and that + attribute is the only thing ``BaseModel.calculate_metrics`` consults + before deciding whether to write anything: ``if not metrics or not + self.run_id`` returns early. So a caller with no run -- a pipeline -- gets + a model that computes nothing into the ``Metric`` table, during training + or after it, without the caller having to intercept anything. + + This is the sole reason ``run_id`` is nullable rather than required. + """ + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + ctx.put_ref("task_name", "ATask") + + _build_unit("ModelA", run_id=None)(ctx) + + assert ctx.get("model").run_id is None + + +def test_a_real_run_id_is_attached_to_the_model(fake_registry): + """The mirror of the above: a run's model has to be able to log.""" + ctx = ExecutionContext() + ctx.put("x", {"train": None, "validation": None}) + ctx.put("y", {"train": None, "validation": None}) + ctx.put("n_labels", None) + ctx.put_ref("task_name", "ATask") + + _build_unit("ModelA", run_id=17)(ctx) + + assert ctx.get("model").run_id == 17 diff --git a/tests/back/units/test_context.py b/tests/back/units/test_context.py index 4aa201240..45bce97a9 100644 --- a/tests/back/units/test_context.py +++ b/tests/back/units/test_context.py @@ -178,3 +178,42 @@ class LiveThing: assert ctx.get("model") is live_object assert ctx.require("model") is live_object + + +def test_origin_says_which_half_a_key_is_in(): + ctx = ExecutionContext(refs={"dataset_path": "/tmp/ds"}) + ctx.put("dataset", object()) + + assert ctx.origin("dataset_path") == "ref" + assert ctx.origin("dataset") == "cache" + assert ctx.origin("nothing_here") is None + + +def test_origin_prefers_the_cache_the_same_way_get_does(): + """One key in both halves has to have one answer, not two. + + ``get`` and ``require`` read the cache first, so anything deciding how to + move a value has to agree with them or the value would be read from one + half and written as if it came from the other. + """ + live = object() + ctx = ExecutionContext(refs={"x": 1}) + ctx.put("x", live) + + assert ctx.origin("x") == "cache" + assert ctx.get("x") is live + assert ctx.require("x") is live + + +def test_origin_does_not_copy_the_reference_half(): + """The reason the method exists: asking must not cost a deep copy. + + ``key in ctx.refs`` answers the same question, but ``refs`` returns a deep + copy of every reference the context holds, so asking once per edge of a + graph would copy the whole reference half once per edge. + """ + ctx = ExecutionContext(refs={"nested": {"deep": [1, 2, 3]}}) + inner = ctx._refs["nested"] + + assert ctx.origin("nested") == "ref" + assert ctx._refs["nested"] is inner diff --git a/tests/back/units/test_evaluate_model_to_artifact_unit.py b/tests/back/units/test_evaluate_model_to_artifact_unit.py new file mode 100644 index 000000000..2e2ff53dc --- /dev/null +++ b/tests/back/units/test_evaluate_model_to_artifact_unit.py @@ -0,0 +1,156 @@ +"""Tests for EvaluateModelToArtifactUnit, and its agreement with its sibling. + +The agreement test is the load-bearing one. Two ways to score the same model +that both pass their own tests and return different numbers is the exact +failure this repository has already had once, when a prediction endpoint kept +its own copy of the three steps a prediction takes. +""" + +import pytest + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.evaluate_model_to_artifact_unit import ( + EvaluateModelToArtifactUnit, +) + + +class _Accuracy: + @staticmethod + def score(y_true, y_pred): + return 0.75 + + +class _AlwaysNaN: + @staticmethod + def score(y_true, y_pred): + return float("nan") + + +class _Model: + """A model with the runtime state ModelFactory would have attached.""" + + def __init__(self, run_id=None, metrics=(_Accuracy,), with_data=True): + self.run_id = run_id + self.train_metrics = list(metrics) + self.validation_metrics = list(metrics) + self.test_metrics = list(metrics) + splits = {"train": "x-train", "validation": "x-val", "test": "x-test"} + self.x_data = splits if with_data else None + self.y_data = splits if with_data else None + self.saved = [] + + def predict(self, x_data): + return [0, 1] + + def prepare_output(self, y_data, is_fit=False): + return [0, 1] + + # The two methods under test come from BaseModel; this stands in for it + # closely enough to exercise the unit without training anything. + compute_metrics = None # replaced below + + +def _model(**kwargs): + """A ``_Model`` borrowing BaseModel's real compute_metrics/_save_metrics.""" + from DashAI.back.models.base_model import BaseModel + + model = _Model(**kwargs) + model.compute_metrics = BaseModel.compute_metrics.__get__(model, _Model) + return model + + +def test_the_metrics_are_published_per_split(): + ctx = ExecutionContext() + ctx.put("model", _model()) + + EvaluateModelToArtifactUnit(splits=["TRAIN", "TEST"])(ctx) + + assert ctx.require("metrics") == { + "train": {"_Accuracy": 0.75}, + "test": {"_Accuracy": 0.75}, + } + + +def test_the_metrics_travel_as_plain_data(): + """A ref, not a cached object: what a caller records is plain numbers.""" + ctx = ExecutionContext() + ctx.put("model", _model()) + + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ctx) + + assert ctx.origin("metrics") == "ref" + + +def test_a_model_with_no_run_is_still_evaluated(): + """The whole point: no Run row, and the numbers still come out. + + ``EvaluateModelUnit`` cannot do this. Its metric rows are keyed by a + foreign key to ``run.id``, so with no run it has nowhere to write, and + ``calculate_metrics`` returns without scoring anything. + """ + ctx = ExecutionContext() + ctx.put("model", _model(run_id=None)) + + EvaluateModelToArtifactUnit(splits=["VALIDATION"])(ctx) + + assert ctx.require("metrics") == {"validation": {"_Accuracy": 0.75}} + + +def test_a_split_with_nothing_to_score_is_left_out(): + """Absent is not the same as scored zero metrics. + + An empty entry would claim the split was evaluated and produced no + metrics, which is a different statement from there being no data for it. + """ + ctx = ExecutionContext() + ctx.put("model", _model(with_data=False)) + + EvaluateModelToArtifactUnit(splits=["TRAIN", "TEST"])(ctx) + + assert ctx.require("metrics") == {} + + +def test_a_non_finite_score_is_dropped_but_the_split_is_still_reported(): + """A split whose every metric was non-finite was still evaluated.""" + ctx = ExecutionContext() + ctx.put("model", _model(metrics=(_AlwaysNaN,))) + + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ctx) + + assert ctx.require("metrics") == {"train": {}} + + +def test_the_unit_needs_a_model(): + with pytest.raises(UnitContractError, match="'model'"): + EvaluateModelToArtifactUnit(splits=["TRAIN"])(ExecutionContext()) + + +def test_the_two_evaluation_paths_agree_on_the_same_model_and_split(monkeypatch): + """The numbers the artifact carries are the numbers a run would have logged. + + Both paths score through ``BaseModel.compute_metrics``, so this fixes that + they cannot drift into two answers for the same model. + """ + from DashAI.back.models.base_model import BaseModel + + logged = {} + + def _capture(self, split, level, results, log_index=None): + logged[split.value] = results + + monkeypatch.setattr(BaseModel, "_save_metrics", _capture) + + # The run path: a model with a run id, going through calculate_metrics. + with_run = _model(run_id=7) + with_run._save_metrics = BaseModel._save_metrics.__get__(with_run, _Model) + BaseModel.calculate_metrics.__get__(with_run, _Model)( + split=SplitEnum.TEST, level=LevelEnum.LAST + ) + + # The artifact path: the same model, no run, through the unit. + ctx = ExecutionContext() + ctx.put("model", _model(run_id=None)) + EvaluateModelToArtifactUnit(splits=["TEST"])(ctx) + + assert logged["test"] == ctx.require("metrics")["test"] diff --git a/tests/back/units/test_evaluate_model_unit.py b/tests/back/units/test_evaluate_model_unit.py index 2f6f8c187..6b4a95b6c 100644 --- a/tests/back/units/test_evaluate_model_unit.py +++ b/tests/back/units/test_evaluate_model_unit.py @@ -2,23 +2,38 @@ import pytest +from DashAI.back.job.base_job import JobError from DashAI.back.units.context import ExecutionContext, UnitContractError from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit def test_call_refuses_to_run_without_a_run_id(): - """Regression: a missing run_id must fail loudly, not silently no-op. + """Regression: a missing run id must fail loudly, not silently no-op. The idempotency check filters an existing-metric query by ``run_id``; a - silently-``None`` value would match no row regardless of what was already - logged, and — if the model were also somehow detached from its run — - ``BaseModel.calculate_metrics`` no-ops on a falsy ``run_id`` too. Before - this fix ``run_id`` was read with ``ctx.get`` and absent from - ``REQUIRES``, so the unit could "succeed" having written zero metrics - instead of surfacing the missing wiring. + ``None`` value would match no row regardless of what was already logged, + and -- if the model were also detached from its run -- + ``BaseModel.calculate_metrics`` no-ops on a falsy ``run_id`` too. Either + way the unit could "succeed" having written zero metrics instead of + surfacing the missing wiring. + + The run id is configuration rather than a context key (no unit publishes + it), so the guard lives in ``validate`` instead of in ``REQUIRES``. """ ctx = ExecutionContext() ctx.put("model", object()) - with pytest.raises(UnitContractError, match="'run_id'"): - EvaluateModelUnit()(ctx) + with pytest.raises(JobError, match="no run id"): + EvaluateModelUnit(run_id=None)(ctx) + + +def test_validate_refuses_a_missing_run_id_before_any_metric_is_computed(): + """``__call__`` validates first, so nothing is written on the way out.""" + with pytest.raises(JobError, match="no run id"): + EvaluateModelUnit(run_id=None).validate(ExecutionContext()) + + +def test_the_unit_still_needs_a_model(): + """``run_id`` left REQUIRES; ``model`` did not.""" + with pytest.raises(UnitContractError, match="'model'"): + EvaluateModelUnit(run_id=1)(ExecutionContext()) diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py index 3ce22551c..b35af3f65 100644 --- a/tests/back/units/test_fit_model_unit.py +++ b/tests/back/units/test_fit_model_unit.py @@ -3,14 +3,34 @@ import pytest from kink import di +from DashAI.back.job.base_job import JobError +from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer from DashAI.back.units.context import ExecutionContext, UnitContractError from DashAI.back.units.fit_model_unit import FitModelUnit - -def _unit(optimizer_name="OptunaOptimizer", goal_metric="Accuracy"): +# Two trials over one hyperparameter: the shape create_plots expects, +# small enough that only the filenames are under test here. +_TRIALS = [ + {"params": {"C": 0.1}, "value": 0.5}, + {"params": {"C": 1.0}, "value": 0.9}, +] +_GOAL_METRIC = {"name": "Accuracy", "metadata": {"maximize": True}} + + +def _unit( + optimizer_name="OptunaOptimizer", + goal_metric="Accuracy", + run_id=None, + artifact_prefix=None, +): + # run_id and artifact_prefix are read straight from the config, with no + # default: omitting one is a KeyError rather than a silently disabled + # runtime-state assertion or a plot filename nobody chose. return FitModelUnit( optimizer={"component": optimizer_name, "params": {}}, goal_metric=goal_metric, + run_id=run_id, + artifact_prefix=artifact_prefix, ) @@ -72,3 +92,74 @@ def __init__(self, **params): assert not ctx.has("goal_metric") finally: del di["component_registry"] + + +def test_the_plot_filenames_come_from_the_artifact_prefix_when_there_is_one(): + """A caller that is not a run names the plots itself. + + Two pipeline executions have no run id to tell them apart, so without a + prefix both would write ``history_objective_plot_None.pickle`` into the + shared runs directory and the second would overwrite the first. + """ + optimizer = OptunaOptimizer() + + filenames, _ = optimizer.create_plots( + _TRIALS, + None, + n_params=1, + goal_metric=_GOAL_METRIC, + artifact_prefix="pipeline-3-fit", + ) + + assert filenames == [ + "history_objective_plot_pipeline-3-fit.pickle", + "slice_plot_pipeline-3-fit.pickle", + ] + + +def test_without_an_artifact_prefix_the_filenames_still_come_from_the_run_id(): + """The default has to leave a real run's filenames exactly as they were.""" + optimizer = OptunaOptimizer() + + named_by_default, _ = optimizer.create_plots( + _TRIALS, 42, n_params=1, goal_metric=_GOAL_METRIC + ) + named_explicitly, _ = optimizer.create_plots( + _TRIALS, 42, n_params=1, goal_metric=_GOAL_METRIC, artifact_prefix=None + ) + + assert named_by_default == [ + "history_objective_plot_42.pickle", + "slice_plot_42.pickle", + ] + assert named_explicitly == named_by_default + + +class _Detached: + """What an optimizer must never return: a model with no data attached.""" + + x_data = None + run_id = None + + +class _Attached: + """What ModelFactory produces: the splits hang off the instance.""" + + x_data = {"train": "x"} + run_id = None + + +def test_a_model_detached_from_its_data_is_refused_even_with_no_run(): + """The guard used to key on run_id, which made it a no-op for a pipeline. + + A pipeline always has run_id None, so returning early on that skipped the + check for the one caller with no other signal: it would finish with an + empty metrics artifact instead of an error. What both callers need is that + the model still carries what scoring reads, so that is what is checked. + """ + with pytest.raises(JobError, match="detached from its data"): + FitModelUnit._assert_model_keeps_its_runtime_state(_Detached()) + + +def test_a_model_that_kept_its_data_passes_without_a_run(): + FitModelUnit._assert_model_keeps_its_runtime_state(_Attached()) diff --git a/tests/back/units/test_save_model_unit.py b/tests/back/units/test_save_model_unit.py new file mode 100644 index 000000000..7aeb1558c --- /dev/null +++ b/tests/back/units/test_save_model_unit.py @@ -0,0 +1,91 @@ +"""Tests for SaveModelUnit's contract, independent of a real model or disk.""" + +import os + +import pytest +from kink import di + +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.save_model_unit import SaveModelUnit + + +class _RecordingModel: + """Stands in for a trained model, recording where it was asked to go.""" + + def __init__(self) -> None: + self.saved_to = None + + def save(self, path) -> None: + self.saved_to = path + + +@pytest.fixture(name="runs_path") +def fixture_runs_path(tmp_path): + di["config"] = {"RUNS_PATH": str(tmp_path)} + yield str(tmp_path) + del di["config"] + + +def test_the_model_lands_in_a_directory_named_by_the_prefix(runs_path): + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + SaveModelUnit(artifact_prefix="pipeline-3-save")(ctx) + + assert model.saved_to == os.path.join(runs_path, "pipeline-3-save") + assert ctx.require("model_path") == model.saved_to + + +def test_a_numeric_run_id_is_a_valid_prefix(runs_path): + """The path a real run writes to has to keep working unchanged. + + ``ModelJob`` passes ``str(run.id)``, and + ``test_model_job_orchestration.py`` asserts the directory is named exactly + that. A guard that rejected it would break every training run. + """ + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + SaveModelUnit(artifact_prefix="17")(ctx) + + assert model.saved_to == os.path.join(runs_path, "17") + + +@pytest.mark.parametrize( + "prefix", + ["../escape", "a/b", "a\b", "", "C:nope", ".", "with space"], +) +def test_a_prefix_that_could_be_read_as_a_path_is_refused(runs_path, prefix): + """A prefix is a directory name, so it must not be able to leave RUNS_PATH. + + ``os.path.join`` with a separator or a parent reference in the prefix + happily produces a destination outside the runs directory, and the model + would be written there without any error. + """ + ctx = ExecutionContext() + model = _RecordingModel() + ctx.put("model", model) + + with pytest.raises(JobError, match="artifact prefix"): + SaveModelUnit(artifact_prefix=prefix)(ctx) + + assert model.saved_to is None + + +def test_validate_runs_before_anything_is_written(runs_path): + """``__call__`` validates first, so a bad prefix never reaches ``save``.""" + unit = SaveModelUnit(artifact_prefix="../escape") + ctx = ExecutionContext() + ctx.put("model", _RecordingModel()) + + with pytest.raises(JobError): + unit.validate(ctx) + + +def test_the_unit_still_needs_a_model(runs_path): + """``run_id`` left REQUIRES; ``model`` did not.""" + with pytest.raises(UnitContractError, match="'model'"): + SaveModelUnit(artifact_prefix="1")(ExecutionContext()) diff --git a/tests/back/units/test_unit_contracts.py b/tests/back/units/test_unit_contracts.py index 7649bdf6e..60bea26ac 100644 --- a/tests/back/units/test_unit_contracts.py +++ b/tests/back/units/test_unit_contracts.py @@ -9,6 +9,7 @@ import ast import pathlib +import re import pytest @@ -69,6 +70,24 @@ def _context_calls(cls, methods): return keys +def _session_names(cls): + """Names bound by ``with session_factory() as :`` inside the class.""" + names = set() + for node in ast.walk(cls): + if not isinstance(node, (ast.With, ast.AsyncWith)): + continue + for item in node.items: + call = item.context_expr + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Name) + and "session" in call.func.id.lower() + and isinstance(item.optional_vars, ast.Name) + ): + names.add(item.optional_vars.id) + return names + + def _parsed_units(): units = [] for path in _unit_modules(): @@ -159,3 +178,241 @@ def test_a_unit_does_not_use_the_context_as_its_own_scratchpad(name, cls): f"{name} writes and reads {sorted(scratch)} without promising it. " "Keep per-instance state on the unit, not in the shared context." ) + + +def test_no_unit_requires_a_key_that_no_unit_provides(): + """Every declared input has to have a declared producer. + + This is what makes a static graph validator decidable. As long as some key + has no producer anywhere in the palette, "nothing supplies this key" is + ambiguous: it could be an edge the user forgot to draw, or a constant the + engine is expected to inject. A validator cannot tell those apart, and the + consequences are not symmetric -- a wrongly injected id is read as a + legitimate value and fails silently. + + ``run_id`` was the only such key: four units required it and none published + it, because in a job it arrived through ``self.kwargs``. It lives in unit + configuration now, where the schema validates it. The rule that keeps it + that way: a key in ``REQUIRES`` is something an upstream unit produces; a + constant the caller chooses is configuration. + """ + required = {} + provided = set() + for name, cls in UNITS: + provided |= _declared(cls, "PROVIDES") + for key in _declared(cls, "REQUIRES"): + required.setdefault(key, []).append(name) + + orphans = {key: names for key, names in required.items() if key not in provided} + assert not orphans, ( + "These keys are required but no unit provides them: " + f"{ {key: sorted(names) for key, names in sorted(orphans.items())} }. " + "A key with no producer makes 'nothing supplies this' ambiguous for the " + "graph validator. Put caller-chosen constants in the unit's config, not " + "in the context." + ) + + +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_write_domain_rows(name, cls): + """Units read the database; jobs and endpoints write it. + + A unit that persists an application entity cannot be reused by a caller + that has no such entity to persist against -- and it is the caller, not the + unit, that owns transaction boundaries and status transitions. Several units + do open a session, always read-only (``db.get`` / ``db.query``); the ones + whose verb is "save" write to disk and publish the path. + + The one place a unit reaches a write is indirect and named: + ``EvaluateModelUnit`` calls ``BaseModel.calculate_metrics``, which persists + through a session of its own. That is why that unit needs a real ``Run`` + row and refuses to run without one. + """ + # Whatever the unit bound its session to, rather than a fixed list of + # names: a unit that writes through ``as s:`` would otherwise slip past. + sessions = {"db", "session", "sess"} | _session_names(cls) + + writes = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in {"add", "add_all", "delete", "commit", "flush"} + and isinstance(node.func.value, ast.Name) + and node.func.value.id in sessions + ): + writes.add(f"{node.func.value.id}.{node.func.attr}") + + assert not writes, ( + f"{name} writes to the database ({sorted(writes)}). Transaction " + "boundaries and status transitions belong to the caller, and a unit " + "that persists an application entity cannot be reused by a caller that " + "has none." + ) + + +# --- The configuration side of the contract --------------------------------- +# +# Until now this file audited only the context. Configuration went unchecked, +# and it could not be checked here: telling a user-facing field from an +# engine-supplied one meant reading a flag off the emitted JSON schema, which +# needs the imported class. Splitting the two declarations apart put both within +# reach of the AST. + +#: Names that usually belong to plumbing: an id, a path, a prefix. Not a +#: definition, a net -- a new field shaped like this has to be classified rather +#: than land in a form by default. +_PLUMBING_SHAPED = re.compile(r"(_id|_path|_prefix)$|^path$") + +#: Fields a user does choose, but that a plain text input cannot render: they +#: need to select an entity, or to type a row. They stay in the schema, because +#: needing a purpose-built widget is a rendering problem and hiding them would +#: answer the wrong question. The value records what each one needs, so the day +#: a node form is built it can tell "do not show this" from "show this +#: differently" instead of rediscovering the distinction. +NEEDS_A_WIDGET = { + ("save_dataset_to_path_unit.py", "path"): "directory picker", + ("load_dataset_unit.py", "dataset_id"): "dataset selector", + ("load_dataset_unit.py", "notebook_id"): "notebook selector", + ("load_datafile_dataset_unit.py", "datafile_id"): "datafile selector", + ("load_run_model_unit.py", "run_id"): "run selector", + ("load_trained_model_unit.py", "run_id"): "run selector", + ("run_exploration_unit.py", "explorer_id"): "exploration selector", + ("save_exploration_unit.py", "explorer_id"): "exploration selector", + ("generate_global_explanation_unit.py", "explainer_id"): "explainer selector", + ("generate_local_explanation_unit.py", "explainer_id"): "explainer selector", + ("generate_local_explanation_unit.py", "instance_dataset_id"): "dataset selector", + ("build_manual_input_unit.py", "manual_input_data"): "typed row editor", + ("generate_local_explanation_unit.py", "manual_input_data"): "typed row editor", +} + + +def _schema_fields(tree): + """Names annotated in the module's ``*Schema`` class, from the AST.""" + names = set() + for node in ast.walk(tree): + if not (isinstance(node, ast.ClassDef) and node.name.endswith("Schema")): + continue + for statement in node.body: + if isinstance(statement, ast.AnnAssign) and isinstance( + statement.target, ast.Name + ): + names.add(statement.target.id) + return names + + +def _config_reads(cls): + """Every ``self.config["key"]`` and ``self.config.get("key")`` in the class.""" + keys = set() + for node in ast.walk(cls): + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Attribute) + and node.value.attr == "config" + and isinstance(node.slice, ast.Constant) + and isinstance(node.slice.value, str) + ): + keys.add(node.slice.value) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "get" + and isinstance(node.func.value, ast.Attribute) + and node.func.value.attr == "config" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + keys.add(node.args[0].value) + return keys + + +#: (filename, unit class node, module tree) for every unit. +UNITS_WITH_TREES = [] +for _path in _unit_modules(): + _tree = ast.parse(_path.read_text(encoding="utf-8")) + _cls = _unit_class(_tree) + if _cls is not None: + UNITS_WITH_TREES.append((_path.name, _cls, _tree)) + + +_IDS = [name for name, _, _ in UNITS_WITH_TREES] + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_every_config_key_a_unit_reads_is_declared(name, cls, tree): + """In the schema if a user fills it in, in RUNTIME_PARAMS if not. + + An undeclared key is a raw ``KeyError`` waiting for the first caller that + builds the unit from its declarations instead of copying an existing call. + """ + read = _config_reads(cls) + declared = _schema_fields(tree) | _declared(cls, "RUNTIME_PARAMS") + + undeclared = read - declared + assert not undeclared, ( + f"{name} reads {sorted(undeclared)} from its configuration without " + "declaring it. Put it in the schema if a user chooses it, or in " + "RUNTIME_PARAMS if whatever runs the unit supplies it." + ) + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_runtime_param_is_not_also_in_the_schema(name, cls, tree): + """The two answers are mutually exclusive. + + A name in both would be exposed to the front and overridden by the caller + at once -- a field the user is invited to fill in and whose value is then + discarded. + """ + overlap = _declared(cls, "RUNTIME_PARAMS") & _schema_fields(tree) + assert not overlap, f"{name} declares {sorted(overlap)} twice." + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_runtime_param_is_actually_read(name, cls, tree): + """A declared name nothing reads is a promise to nobody.""" + unread = _declared(cls, "RUNTIME_PARAMS") - _config_reads(cls) + assert not unread, ( + f"{name} declares {sorted(unread)} in RUNTIME_PARAMS but never reads it." + ) + + +@pytest.mark.parametrize(("name", "cls", "tree"), UNITS_WITH_TREES, ids=_IDS) +def test_a_plumbing_shaped_field_is_classified_one_way_or_the_other(name, cls, tree): + """A new field named like plumbing cannot default into a form. + + The net has a known hole: ``session_splits`` and + ``trust_inherited_metadata`` are runtime params and match nothing, which is + why the exact set in ``tests/back/api/test_units_api.py`` exists as well. + What this catches is the common case -- a new ``*_id`` or ``*_path`` added + without anyone deciding what it is. + """ + runtime = _declared(cls, "RUNTIME_PARAMS") + + unclassified = [ + field + for field in _schema_fields(tree) + if _PLUMBING_SHAPED.search(field) + and field not in runtime + and (name, field) not in NEEDS_A_WIDGET + ] + + assert not unclassified, ( + f"{name} has schema fields named like plumbing that nobody classified: " + f"{sorted(unclassified)}. Move them to RUNTIME_PARAMS if whatever runs " + "the unit supplies them, or add them to NEEDS_A_WIDGET with the widget " + "they need." + ) + + +def test_every_entry_in_the_widget_registry_still_exists(): + """Guards the registry itself: a stale entry would make it decorative.""" + declared = { + (name, field) + for name, _, tree in UNITS_WITH_TREES + for field in _schema_fields(tree) + } + + missing = sorted(set(NEEDS_A_WIDGET) - declared) + assert not missing, missing From 46024af00f0d8fdf2bc956d94b32613b9daa9afa Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 16:21:32 -0300 Subject: [PATCH 16/28] Split the preparing unit in two, over the registered splitters PrepareAndSplitUnit carried the partitioning policy inside itself, reached through prepare_for_model_session, and took its configuration as an untyped `splits` dictionary -- the kind of field the atomization notes admit only because there was nothing better to declare. develop meanwhile grew ten splitters, each a registered component with its own schema, multilingual labels and compatibility per task. Those are the better thing: the unit now picks one with a component field, the same way BuildModelUnit picks a model and FitModelUnit picks an optimizer. Two units rather than one with a flag, for two reasons that both bite: A component field carries a single `parent` and the front reads it directly off the property, so a field offering both families would leave the user without a selector at all -- the same wall that made the two explainer units siblings. What comes back has a different *type*. A holdout splitter returns one DatasetDict per side; a fold splitter returns a list of them, plus a trailing entry that is not a fold. Publishing that list as `x` would give one key two shapes, which a contract comparing key names cannot express: a graph would validate and then fail at run time, or quietly train on the wrong thing. So PrepareAndFoldUnit publishes `x_folds` and `y_folds`. The two families are told apart with no renaming and no registry change: component_parent matches any ancestor by name, and the hierarchy already partitions the ten exactly -- PartitionSplitter covers the two holdout splitters, FoldSplitter the eight fold ones. The shared body lives in splitter_scope.py, next to converter_scope.py and for the same reason: one implementation of resolving the task, preparing the dataset and selecting the columns, so the siblings cannot drift into two answers for the same dataset. It takes and returns plain values and never touches the context -- a ctx.put hidden in a helper is invisible to the audit that parses each unit's own source, so a broken PROVIDES would pass it. Two details worth naming: BaseSplitter.__init__ takes a single `splits_data` mapping rather than keyword arguments, so this is the one component field in the units that is not expanded with **params. The instance state is declared on each unit and not in the mixin's __init__. BaseUnit.__init__ comes first in the MRO and does not chain, so a mixin __init__ never runs -- which surfaced as the resolved task being missing rather than as anything about construction. ApplyConverterUnit already does it this way. The splitter's own refusal passes through undecorated: it already names the numbers that explain it, and the caller that knows which run this was frames it from outside, which is how the message the user reads is built today. Contract tests build the context by hand rather than going through a job, including that two of these units in one context do not share a resolved task. The spike is untouched in substance: it only ever used this unit for static validation, which reads REQUIRES and PROVIDES and never constructs anything. 969 passed in units, dag, spike and api; the one failure is the CV-aware explanation indexes still to be ported, which is a later slice. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/initial_components.py | 2 + DashAI/back/units/prepare_and_fold_unit.py | 73 +++++ DashAI/back/units/prepare_and_split_unit.py | 202 +++--------- DashAI/back/units/splitter_scope.py | 303 ++++++++++++++++++ tests/back/api/test_model_job_as_a_graph.py | 24 +- tests/back/api/test_units_api.py | 24 +- .../pipeline_spike/test_dag_engine_spike.py | 2 +- tests/back/units/test_splitting_units.py | 265 +++++++++++++++ 8 files changed, 721 insertions(+), 174 deletions(-) create mode 100644 DashAI/back/units/prepare_and_fold_unit.py create mode 100644 DashAI/back/units/splitter_scope.py create mode 100644 tests/back/units/test_splitting_units.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index e13aeeaac..2339426ab 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -539,6 +539,7 @@ from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit from DashAI.back.units.load_uploaded_dataset_unit import LoadUploadedDatasetUnit from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit 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 @@ -728,6 +729,7 @@ def get_initial_components(): # Units LoadDatasetUnit, PrepareAndSplitUnit, + PrepareAndFoldUnit, BuildModelUnit, FitModelUnit, EvaluateModelUnit, diff --git a/DashAI/back/units/prepare_and_fold_unit.py b/DashAI/back/units/prepare_and_fold_unit.py new file mode 100644 index 000000000..c71842328 --- /dev/null +++ b/DashAI/back/units/prepare_and_fold_unit.py @@ -0,0 +1,73 @@ +"""Unit that prepares a dataset for a task and carves it into folds.""" + +import logging + +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.splitter_scope import ( + SplitterScopeMixin, + fold_splitter_field, + input_columns_field, + output_columns_field, + task_name_field, +) + +log = logging.getLogger(__name__) + + +class PrepareAndFoldSchema(BaseSchema): + task_name: task_name_field() # type: ignore + input_columns: input_columns_field() # type: ignore + output_columns: output_columns_field() # type: ignore + splitter: fold_splitter_field() # type: ignore + + +class PrepareAndFoldUnit(BaseUnit, SplitterScopeMixin): + """Validate a dataset against a task and carve it into cross-validation folds. + + The sibling of ``PrepareAndSplitUnit``, sharing its whole body. The two + differ in which family of splitters they offer and, because of that, in the + shape of what they publish. + + **What comes back is a list, and that is why the keys are different.** A + fold splitter returns one entry per fold plus a trailing entry that is not + a fold: its train partition is every row the folds could use and its test + partition holds the rows reserved for scoring the model that gets kept, + empty when the session reserved none. Publishing that list as ``x`` would + give one key two possible types, which a contract comparing key names + cannot express -- a graph would validate and then fail at run time, or + worse, quietly train on the wrong thing. So it is ``x_folds`` and + ``y_folds``. + + ``split_indexes`` keeps its name because it keeps its meaning -- the rows + of every partition this run produced -- even though a fold payload is + shaped differently from a holdout one. Which shape it is, is answered by + asking the splitter that produced it, not by inspecting the payload. + """ + + SCHEMA = PrepareAndFoldSchema + + REQUIRES = ("dataset", "dataset_id") + PROVIDES = ("x_folds", "y_folds", "n_labels", "task", "split_indexes", "task_name") + + def __init__(self, **config) -> None: + super().__init__(**config) + # Kept on the instance, never in the context: two of these units in one + # context would otherwise overwrite each other's resolved task. + self._task = None + self._splitter_class = None + + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") + dataset_id = ctx.require("dataset_id") + + task, n_labels, x, y = self._prepare(dataset, dataset_id) + x_folds, y_folds, split_indexes = self._split(x, y) + + ctx.put_ref("task_name", self.config["task_name"]) + ctx.put_ref("split_indexes", split_indexes) + ctx.put("task", task) + ctx.put("n_labels", n_labels) + ctx.put("x_folds", x_folds) + ctx.put("y_folds", y_folds) diff --git a/DashAI/back/units/prepare_and_split_unit.py b/DashAI/back/units/prepare_and_split_unit.py index 17df6c52c..26aeb9caf 100644 --- a/DashAI/back/units/prepare_and_split_unit.py +++ b/DashAI/back/units/prepare_and_split_unit.py @@ -1,110 +1,48 @@ """Unit that prepares a dataset for a task and splits it into train/val/test.""" 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.core.schema_fields import BaseSchema 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 +from DashAI.back.units.splitter_scope import ( + SplitterScopeMixin, + input_columns_field, + output_columns_field, + partition_splitter_field, + task_name_field, +) log = logging.getLogger(__name__) class PrepareAndSplitSchema(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: schema_field( - list_field(string_field(), min_items=1), - placeholder=[], - 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="用作模型输入的列名。", - ), - 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 has to predict.", - es="Nombres de las columnas que el modelo debe predecir.", - pt="Nomes das colunas que o modelo deve prever.", - de="Namen der Spalten, die das Modell vorhersagen soll.", - zh="模型需要预测的列名。", - ), - alias=MultilingualString( - en="Output columns", - es="Columnas de salida", - pt="Colunas de saída", - de="Ausgabespalten", - zh="输出列", - ), - ) # type: ignore - splits: schema_field( - dict, - placeholder={ - "splitType": "random", - "train": 0.7, - "test": 0.1, - "validation": 0.2, - }, - description=MultilingualString( - en="Split configuration: a split type plus either train/test/" - "validation index lists or proportions.", - es="Configuración de partición: un tipo de partición y listas de " - "índices o proporciones para entrenamiento/prueba/validación.", - pt="Configuração de divisão: um tipo de divisão e listas de " - "índices ou proporções para treino/teste/validação.", - de="Split-Konfiguration: ein Split-Typ sowie entweder Index-" - "Listen oder Anteile für Training/Test/Validierung.", - zh="划分配置：划分类型，以及训练/测试/验证的索引列表或比例。", - ), - alias=MultilingualString( - en="Splits", - es="Particiones", - pt="Partições", - de="Teilmengen", - zh="数据划分", - ), - ) # type: ignore - - -class PrepareAndSplitUnit(BaseUnit): - """Validate a dataset against a task and split it into train/val/test. - - Runs the task's own validation, counts the labels, applies the requested - split configuration and separates features from targets. + task_name: task_name_field() # type: ignore + input_columns: input_columns_field() # type: ignore + output_columns: output_columns_field() # type: ignore + splitter: partition_splitter_field() # type: ignore + + +class PrepareAndSplitUnit(BaseUnit, SplitterScopeMixin): + """Validate a dataset against a task and split it once into three partitions. + + Runs the task's own validation, counts the labels, separates features from + targets and hands the pair to a holdout splitter. + + The split configuration is a component field rather than an untyped + dictionary: which partitions exist, what they are called and what may be + configured are the splitter's own schema, so the front renders that form + instead of a free-form object nobody could validate. + + The sibling for cross-validation is ``PrepareAndFoldUnit``. They are two + units and not one with a flag for two reasons that both bite. A component + field carries a single ``parent`` and the front reads it directly off the + property, so offering two families from one field leaves the user with no + selector at all. And what comes back has a different *type* -- one + ``DatasetDict`` against a list of them -- which a contract that compares + key names cannot express: the same key holding two shapes would validate + statically and fail at run time. The fold unit therefore publishes + ``x_folds`` and ``y_folds`` rather than reusing ``x`` and ``y``. """ SCHEMA = PrepareAndSplitSchema @@ -115,69 +53,21 @@ class PrepareAndSplitUnit(BaseUnit): REQUIRES = ("dataset", "dataset_id") PROVIDES = ("x", "y", "n_labels", "task", "split_indexes", "task_name") - def execute(self, ctx: ExecutionContext) -> None: - from kink import di - - from DashAI.back.dataloaders.classes.dashai_dataset import ( - prepare_for_model_session, - select_columns, - split_dataset, - ) + def __init__(self, **config) -> None: + super().__init__(**config) + # Kept on the instance, never in the context: two of these units in one + # context would otherwise overwrite each other's resolved task. + self._task = None + self._splitter_class = None - component_registry = di["component_registry"] - - task_name: str = self.config["task_name"] - input_columns = self.config["input_columns"] - output_columns = self.config["output_columns"] - splits = self.config["splits"] - - loaded_dataset = ctx.require("dataset") + def execute(self, ctx: ExecutionContext) -> None: + dataset = ctx.require("dataset") dataset_id = ctx.require("dataset_id") - 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 - - try: - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=input_columns, - output_columns=output_columns, - ) - n_labels = task.num_labels(prepared_dataset, output_columns[0]) - - prepared_dataset, splits = prepare_for_model_session( - dataset=prepared_dataset, - splits=splits, - output_columns=output_columns, - ) - - split_indexes = { - "train_indexes": splits["train_indexes"], - "test_indexes": splits["test_indexes"], - "val_indexes": splits["val_indexes"], - } - - x, y = select_columns( - prepared_dataset, - input_columns, - output_columns, - ) - - x = split_dataset(x) - y = split_dataset(y) - - except Exception as e: - log.exception(e) - raise JobError( - f"Can not prepare Dataset {dataset_id} for Task {task_name}", - ) from e - - ctx.put_ref("task_name", task_name) + task, n_labels, x, y = self._prepare(dataset, dataset_id) + x, y, split_indexes = self._split(x, y) + + ctx.put_ref("task_name", self.config["task_name"]) ctx.put_ref("split_indexes", split_indexes) ctx.put("task", task) ctx.put("n_labels", n_labels) diff --git a/DashAI/back/units/splitter_scope.py b/DashAI/back/units/splitter_scope.py new file mode 100644 index 000000000..1954ba823 --- /dev/null +++ b/DashAI/back/units/splitter_scope.py @@ -0,0 +1,303 @@ +"""Shared body of the two units that prepare a dataset and partition it. + +``PrepareAndSplitUnit`` and ``PrepareAndFoldUnit`` do the same four things -- +resolve the task, prepare the dataset for it, separate features from targets, +and hand the pair to a splitter -- and differ only in which family of splitters +they offer and in the shape of what comes back. Keeping the body here is what +stops the two from drifting into two answers for the same dataset. + +Named ``SplitterScopeMixin`` rather than ``BaseSomething`` on purpose: the +registry derives a component's type by walking its ``__mro__`` for a class whose +name contains "Base" and that declares ``TYPE``, and demands exactly one. A +shared parent called ``Base*`` would be a second candidate and would break the +registration of every unit that inherited it. +""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Tuple + +from DashAI.back.core.schema_fields import ( + component_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + from DashAI.back.splitters.base_splitter import BaseSplitter + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +def task_name_field(): + return 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="任务" + ), + ) + + +def input_columns_field(): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + 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="用作模型输入的列名。", + ), + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + ) + + +def output_columns_field(): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns the model has to predict.", + es="Nombres de las columnas que el modelo debe predecir.", + pt="Nomes das colunas que o modelo deve prever.", + de="Namen der Spalten, die das Modell vorhersagen soll.", + zh="模型需要预测的列名。", + ), + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + ) + + +def _splitter_field(parent: str, placeholder: Dict[str, Any]): + """A splitter chosen from one family, with its own configuration. + + Parameters + ---------- + parent : str + Class name every offered splitter has in its ``__mro__``. The two + families are told apart here rather than by a flag, because the front + resolves the choices with ``component_parent``, which matches any + ancestor by name: ``PartitionSplitter`` offers the two holdout + splitters and ``FoldSplitter`` the eight fold ones, with no renaming + and no change to the registry. + placeholder : dict + The ``{"component": …, "params": {…}}`` value the form starts on. + """ + return schema_field( + component_field(parent=parent), + placeholder=placeholder, + description=MultilingualString( + en="Splitter that decides how the dataset is partitioned, along " + "with its own configuration.", + es="Particionador que decide cómo se divide el conjunto de datos, " + "junto con su propia configuración.", + pt="Divisor que decide como o conjunto de dados é particionado, " + "junto com a sua própria configuração.", + de="Splitter, der über die Aufteilung des Datensatzes entscheidet, " + "samt seiner eigenen Konfiguration.", + zh="决定数据集如何划分的划分器及其自身配置。", + ), + alias=MultilingualString( + en="Splitter", + es="Particionador", + pt="Divisor", + de="Splitter", + zh="划分器", + ), + ) + + +def partition_splitter_field(): + """The field of the unit that splits once, offering the holdout family.""" + return _splitter_field( + parent="PartitionSplitter", + placeholder={ + "component": "HoldoutSplitter", + "params": { + "train": 0.6, + "test": 0.2, + "validation": 0.2, + "stratify": False, + "shuffle": True, + "random_state": 42, + }, + }, + ) + + +def fold_splitter_field(): + """The field of the unit that splits into folds, offering the fold family.""" + return _splitter_field( + parent="FoldSplitter", + placeholder={ + "component": "KFoldSplitter", + "params": { + "n_splits": 5, + "test_size": 0.1, + "shuffle": True, + "random_state": 42, + }, + }, + ) + + +class SplitterScopeMixin: + """Prepare a dataset for a task and hand it to a splitter. + + Every method here takes and returns plain values and never touches the + execution context. A ``ctx.put`` hidden in a shared helper is invisible to + the audit that parses each unit's own source, so a broken ``PROVIDES`` + would pass it. + """ + + #: Declared here so both the memoizing helpers and a reader can see what + #: state an instance carries; each unit sets them in its own ``__init__``, + #: because ``BaseUnit.__init__`` comes first in the MRO and does not chain. + _task = None + _splitter_class = None + + @property + def splitter_name(self) -> str: + return self.config["splitter"]["component"] + + @property + def splitter_params(self) -> Dict[str, Any]: + return self.config["splitter"]["params"] + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task, memoized on this unit. + + On the instance rather than in the context: a context can hold two of + these units, and a context-global cache would silently give the second + one the first one's task. + """ + if self._task is not None: + return self._task + + from kink import di + + task_name: str = self.config["task_name"] + try: + self._task = di["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 + return self._task + + def _resolve_splitter(self) -> "BaseSplitter": + """Build the configured splitter. + + ``BaseSplitter.__init__`` takes a single ``splits_data`` mapping rather + than keyword arguments, so this is the one component field in the units + that is not expanded with ``**params``. The shape is the splitter's own + schema either way; only how it is handed over differs. + """ + from kink import di + + splitter_name = self.splitter_name + if self._splitter_class is None: + try: + self._splitter_class = di["component_registry"][splitter_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Splitter with name {splitter_name} in registry.", + ) from e + + try: + return self._splitter_class(splits_data=dict(self.splitter_params)) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instantiating splitter {splitter_name}, {e}", + ) from e + + def _prepare( + self, dataset: "DashAIDataset", dataset_id: Any + ) -> Tuple["BaseTask", int, "DashAIDataset", "DashAIDataset"]: + """Validate the dataset against the task and separate x from y. + + Returns + ------- + tuple + The task, the number of labels, and the input and output datasets, + in that order. Nothing here is written to the context: the unit + that called this decides what it promises. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import select_columns + + task = self._resolve_task() + task_name: str = self.config["task_name"] + input_columns: List[str] = self.config["input_columns"] + output_columns: List[str] = self.config["output_columns"] + + try: + prepared_dataset = task.prepare_for_task( + dataset=dataset, + input_columns=input_columns, + output_columns=output_columns, + ) + n_labels = task.num_labels(prepared_dataset, output_columns[0]) + except Exception as e: + log.exception(e) + raise JobError( + f"Can not prepare Dataset {dataset_id} for Task {task_name}", + ) from e + + try: + # Read from the prepared dataset rather than the loaded one: a task + # may reorder the rows, and forecasting does, sorting them by date + # so a temporal splitter carves real periods of time. Selecting + # from the loaded dataset would drop that work on the floor. + x, y = select_columns(prepared_dataset, input_columns, output_columns) + except Exception as e: + log.exception(e) + raise JobError( + f"Error selecting input and output columns from dataset {dataset_id}", + ) from e + + return task, n_labels, x, y + + def _split(self, x: "DashAIDataset", y: "DashAIDataset"): + """Partition the pair with the configured splitter. + + The splitter's own complaint is passed through as the whole message, + undecorated. It already names the numbers that explain the refusal -- + how many folds against how many rows -- and a caller that wants to say + which run this was frames it from outside, which is how the message + the user reads is built today. + """ + splitter = self._resolve_splitter() + try: + return splitter.split(x, y) + except JobError: + raise + except Exception as e: + log.exception(e) + raise JobError(str(e)) from e diff --git a/tests/back/api/test_model_job_as_a_graph.py b/tests/back/api/test_model_job_as_a_graph.py index a8ce2473c..8e2794735 100644 --- a/tests/back/api/test_model_job_as_a_graph.py +++ b/tests/back/api/test_model_job_as_a_graph.py @@ -44,6 +44,7 @@ from DashAI.back.metrics.base_metric import BaseMetric from DashAI.back.models.base_model import BaseModel from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.splitters.holdout import HoldoutSplitter from DashAI.back.tasks.base_task import BaseTask from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.evaluate_model_to_artifact_unit import ( @@ -54,16 +55,16 @@ from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit from DashAI.back.units.save_model_unit import SaveModelUnit -SPLITS = { - "train": 0.5, - "test": 0.2, - "validation": 0.3, - "is_random": True, - "has_changed": True, - "seed": 42, - "shuffle": True, - "stratify": False, - "splitType": "random", +SPLITTER = { + "component": "HoldoutSplitter", + "params": { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "shuffle": True, + "stratify": False, + "random_state": 42, + }, } @@ -151,6 +152,7 @@ def setup_graph_registry(client): PipelineJob, LoadDatasetUnit, PrepareAndSplitUnit, + HoldoutSplitter, BuildModelUnit, FitModelUnit, EvaluateModelToArtifactUnit, @@ -187,7 +189,7 @@ def _model_job_blocks(dataset_id: int): "task_name": "GraphTask", "input_columns": ["SepalLengthCm", "SepalWidthCm"], "output_columns": ["Species"], - "splits": SPLITS, + "splitter": SPLITTER, }, } ], diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index c965eb3d9..fea24053a 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -6,6 +6,7 @@ EXPECTED_UNITS = { "LoadDatasetUnit", "PrepareAndSplitUnit", + "PrepareAndFoldUnit", "BuildModelUnit", "FitModelUnit", "EvaluateModelUnit", @@ -64,12 +65,23 @@ def test_unit_schemas_describe_their_configuration(units): "dataset_id", "notebook_id", } - assert set(units["PrepareAndSplitUnit"]["schema"]["properties"]) == { - "task_name", - "input_columns", - "output_columns", - "splits", - } + # The two splitting units are the same form with a different family of + # splitters offered, which is the whole of what separates them. + for name in ("PrepareAndSplitUnit", "PrepareAndFoldUnit"): + assert set(units[name]["schema"]["properties"]) == { + "task_name", + "input_columns", + "output_columns", + "splitter", + }, name + assert ( + units["PrepareAndSplitUnit"]["schema"]["properties"]["splitter"]["parent"] + == "PartitionSplitter" + ) + assert ( + units["PrepareAndFoldUnit"]["schema"]["properties"]["splitter"]["parent"] + == "FoldSplitter" + ) assert "model" in units["BuildModelUnit"]["schema"]["properties"] assert "optimizer" in units["FitModelUnit"]["schema"]["properties"] assert set(units["ApplyConverterUnit"]["schema"]["properties"]) == { diff --git a/tests/back/pipeline_spike/test_dag_engine_spike.py b/tests/back/pipeline_spike/test_dag_engine_spike.py index 797a88968..19600ec55 100644 --- a/tests/back/pipeline_spike/test_dag_engine_spike.py +++ b/tests/back/pipeline_spike/test_dag_engine_spike.py @@ -303,7 +303,7 @@ def test_a_key_a_middle_node_does_not_republish_needs_an_edge_around_it(): "apply", ApplyConverterUnit(converter=_MIN_MAX, scope=FULL_SCOPE, target=None), ) - split = Node("split", PrepareAndSplitUnit(splits={})) + split = Node("split", PrepareAndSplitUnit(splitter={})) bundled = Graph( [load, apply_, split], [*connect(load, apply_), *connect(apply_, split)] diff --git a/tests/back/units/test_splitting_units.py b/tests/back/units/test_splitting_units.py new file mode 100644 index 000000000..e403db6fd --- /dev/null +++ b/tests/back/units/test_splitting_units.py @@ -0,0 +1,265 @@ +"""Contract tests for the two units that prepare a dataset and partition it. + +Built on a hand-made ``ExecutionContext`` rather than through a job: the point +is what each unit reads, promises and refuses on its own, which a job that +always wires the context correctly cannot show. +""" + +import pandas as pd +import pyarrow as pa +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.splitters.holdout import HoldoutSplitter +from DashAI.back.splitters.k_fold import KFoldSplitter +from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.types.value_types import Float +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit + +ROWS = 20 + + +class SplitTask(BaseTask): + name: str = "SplitTask" + metadata: dict = { + "inputs_types": [], + "outputs_types": [], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def num_labels(self, dataset, output_column): + return 2 + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "SplitTask": {"class": SplitTask}, + "HoldoutSplitter": {"class": HoldoutSplitter}, + "KFoldSplitter": {"class": KFoldSplitter}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _dataset(): + frame = pd.DataFrame( + { + "a": [float(i) for i in range(ROWS)], + "b": [float(i % 3) for i in range(ROWS)], + "y": [float(i % 2) for i in range(ROWS)], + } + ) + types = {name: Float(arrow_type=pa.float64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +def _context(): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put_ref("dataset_id", 7) + return ctx + + +def _split_unit(**splitter_params): + params = { + "train": 0.5, + "test": 0.25, + "validation": 0.25, + "shuffle": True, + "stratify": False, + "random_state": 42, + } + params.update(splitter_params) + return PrepareAndSplitUnit( + task_name="SplitTask", + input_columns=["a", "b"], + output_columns=["y"], + splitter={"component": "HoldoutSplitter", "params": params}, + ) + + +def _fold_unit(**splitter_params): + params = { + "n_splits": 4, + "test_size": 0.25, + "shuffle": True, + "random_state": 42, + } + params.update(splitter_params) + return PrepareAndFoldUnit( + task_name="SplitTask", + input_columns=["a", "b"], + output_columns=["y"], + splitter={"component": "KFoldSplitter", "params": params}, + ) + + +# --------------------------------------------------------------------------- # +# What each unit promises +# --------------------------------------------------------------------------- # + + +def test_the_holdout_unit_publishes_one_dataset_dict_per_side(registry): + ctx = _context() + + _split_unit()(ctx) + + assert set(ctx.require("x")) == {"train", "test", "validation"} + assert set(ctx.require("y")) == {"train", "test", "validation"} + assert set(ctx.require("split_indexes")) == { + "train_indexes", + "test_indexes", + "val_indexes", + } + + +def test_the_fold_unit_publishes_a_list_and_not_the_holdout_keys(registry): + """The reason the two are separate units and not one with a flag. + + Reusing ``x`` for a list would give one key two possible types, which a + contract that compares key names cannot express. + """ + ctx = _context() + + _fold_unit()(ctx) + + x_folds = ctx.require("x_folds") + assert isinstance(x_folds, list) + # Four folds plus the trailing entry, which is not a fold. + assert len(x_folds) == 5 + assert not ctx.has("x") + assert not ctx.has("y") + + +def test_every_fold_holds_a_train_and_a_validation_partition(registry): + ctx = _context() + + _fold_unit()(ctx) + + x_folds = ctx.require("x_folds") + for fold in x_folds[:-1]: + assert set(fold) == {"train", "validation"} + # The trailing entry fits the kept model and is scored on the reserved rows. + assert set(x_folds[-1]) == {"train", "test"} + + +def test_the_fold_unit_keeps_the_reserved_rows_out_of_every_fold(registry): + """Checked on the indexes the unit itself published.""" + ctx = _context() + + _fold_unit()(ctx) + + split_indexes = ctx.require("split_indexes") + reserved = set(split_indexes["full_dataset"]["test_indexes"]) + assert reserved + + for name, partitions in split_indexes.items(): + if name == "full_dataset": + continue + seen = set(partitions["train_indexes"]) | set(partitions["validation_indexes"]) + assert reserved.isdisjoint(seen), name + + +def test_a_session_that_reserves_nothing_still_produces_folds(registry): + ctx = _context() + + _fold_unit(test_size=0)(ctx) + + assert ctx.require("split_indexes")["full_dataset"]["test_indexes"] == [] + assert len(ctx.require("x_folds")) == 5 + + +# --------------------------------------------------------------------------- # +# What the two share +# --------------------------------------------------------------------------- # + + +def test_the_two_units_prepare_the_dataset_the_same_way(registry): + """Both run the same body, so the parts before the split cannot diverge. + + The one thing that differs is what the splitter returns, and everything + upstream of it -- the task, the label count, the columns selected -- has to + match or the two would be scoring different datasets. + """ + holdout_ctx, fold_ctx = _context(), _context() + + _split_unit()(holdout_ctx) + _fold_unit()(fold_ctx) + + assert holdout_ctx.require("n_labels") == fold_ctx.require("n_labels") + assert holdout_ctx.require("task_name") == fold_ctx.require("task_name") + assert type(holdout_ctx.require("task")) is type(fold_ctx.require("task")) + + holdout_columns = holdout_ctx.require("x")["train"].column_names + fold_columns = fold_ctx.require("x_folds")[0]["train"].column_names + assert holdout_columns == fold_columns == ["a", "b"] + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_a_unit_refuses_to_run_before_a_dataset_was_loaded(registry, build): + """``__call__`` checks REQUIRES, so the failure names the missing key.""" + with pytest.raises(UnitContractError): + build()(ExecutionContext()) + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_an_unknown_task_is_reported_by_name(registry, build): + unit = build() + unit.config["task_name"] = "ThereIsNoSuchTask" + + with pytest.raises(JobError, match="Unable to find Task with name"): + unit(_context()) + + +@pytest.mark.parametrize("build", [_split_unit, _fold_unit]) +def test_an_unknown_splitter_is_reported_by_name(registry, build): + unit = build() + unit.config["splitter"] = {"component": "ThereIsNoSuchSplitter", "params": {}} + + with pytest.raises(JobError, match="Unable to find Splitter with name"): + unit(_context()) + + +def test_the_splitters_own_refusal_is_passed_through_undecorated(registry): + """More folds than rows is the splitter's diagnosis, not the unit's. + + It already names the numbers that explain the refusal, and a caller that + wants to say which run this was frames it from outside. Decorating it here + would push the numbers into the middle of someone else's sentence. + """ + with pytest.raises(JobError) as error: + _fold_unit(n_splits=ROWS + 1, test_size=0)(_context()) + + message = " ".join(str(error.value).split()) + assert message == ( + f"Number of splits (n_splits={ROWS + 1}) cannot be greater " + f"than the number of samples ({ROWS})." + ) + + +def test_two_units_in_one_context_do_not_share_their_resolved_task(registry): + """Instance state stays on the instance, which is what lets a graph hold two. + + A context-global cache would give the second unit the first one's task, + and the failure would be silent: the wrong task prepares the dataset + without complaining. + """ + first, second = _split_unit(), _fold_unit() + ctx = _context() + + first(ctx) + second(ctx) + + assert first._task is not None + assert second._task is not None + assert first._task is not second._task From e2152858ef176355b3eb8924f2ffb1f0c22ca9c2 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 16:33:30 -0300 Subject: [PATCH 17/28] Let the fit decide what data it sees, and whether validation is part of it Two changes, both policy rather than shape, and both measured by the cross-validation net before they were made. BuildModelUnit no longer takes the data. ModelFactory attached the splits to the model instance at construction, which worked only because a model was fitted once on one split. Fitted over folds it sees different data on every iteration, so binding one partition at build time would leave the metrics describing whichever fold happened to be built with. The unit now needs only the label count -- a property of the dataset, not of a split -- and whoever fits the model points it at what it is being fitted on. REQUIRES loses `x` and `y`, which is a relaxation: every graph that fed it still validates, with two fewer wires. The measurement in the graph test moves from fifteen to thirteen and says why. FitModelUnit gained `validation_during_fit`. It handed the validation partition to `train` unconditionally, and both halves of that are wrong for folds: A model uses validation data to watch the fit and stop early, which is what an ordinary holdout run wants and exactly what a fold does not -- a fold is scored on the rows it held back, so letting the fit watch them measures it on data it was allowed to see. Nothing raises; the score just comes out better than the model deserves. The net recorded four fits in a cross-validated run and none of them receiving validation data, which is the behaviour this field now expresses. The trailing entry a fold splitter produces has no validation partition at all -- it holds the pooled rows and the reserved ones -- so reading x["validation"] is a plain KeyError on the very partition set that fits the model which gets kept. Read with `.get` and the schema's own placeholder, the way the other units read a declared optional field: a caller that builds this unit by hand should not have to name a policy it is happy to leave alone. Also moved the runs directory out of the top of execute and into the branch that needs it. It is only used to name the plots a search produces, so a fit without a search had been requiring a service of its caller for nothing -- which is what made these tests need a container before they could watch a fit. 973 passed across units, dag, spike and api; the one failure is the CV-aware explanation indexes still to be ported. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/units/build_model_unit.py | 28 ++++--- DashAI/back/units/fit_model_unit.py | 84 +++++++++++++++++-- tests/back/api/test_model_job_as_a_graph.py | 13 ++- tests/back/units/test_fit_model_unit.py | 92 +++++++++++++++++++++ 4 files changed, 194 insertions(+), 23 deletions(-) diff --git a/DashAI/back/units/build_model_unit.py b/DashAI/back/units/build_model_unit.py index 72336ad79..e0f5bc0b4 100644 --- a/DashAI/back/units/build_model_unit.py +++ b/DashAI/back/units/build_model_unit.py @@ -100,11 +100,19 @@ class BuildModelSchema(BaseSchema): class BuildModelUnit(BaseUnit): """Instantiate an untrained model bound to its data and metrics. - ``ModelFactory`` attaches the run id, the data splits and the metric - classes to the model instance, which is what later lets the model log - metrics on its own during and after training. The metrics are configured - here rather than in the evaluation unit because models use them *while* - training to log at the step and epoch levels. + ``ModelFactory`` attaches the run id and the metric classes to the model + instance, which is what later lets the model log metrics on its own during + and after training. The metrics are configured here rather than in the + evaluation unit because models use them *while* training to log at the step + and epoch levels. + + **The data is not attached here.** It used to be, and that only worked + because a model was fitted once on one split. A model fitted over folds + sees different data on every iteration, so binding one partition at + construction would leave the metrics describing whichever fold happened to + be built with. Whoever fits the model points it at the data it is being + fitted on, and this unit no longer needs ``x`` or ``y`` at all -- only the + label count, which is a property of the dataset rather than of a split. ``validate`` checks that the model and every component nested in its parameters have been downloaded, so an impossible run is rejected before @@ -129,7 +137,7 @@ class BuildModelUnit(BaseUnit): # upstream could ever satisfy it. It is read without a default on purpose # — a run_id nobody passed would read as "this model has no run", and a # model with no run logs no metrics at all (see BaseModel). - REQUIRES = ("x", "y", "n_labels", "task_name") + REQUIRES = ("n_labels", "task_name") PROVIDES = ("model", "factory", "optimizable_parameters", "model_parameters") RUNTIME_PARAMS = ("run_id",) @@ -232,11 +240,9 @@ def execute(self, ctx: ExecutionContext) -> None: model_class, parameters, run_id, - ctx.require("x"), - ctx.require("y"), - train_metrics, - validation_metrics, - test_metrics, + train_metrics=train_metrics, + validation_metrics=validation_metrics, + test_metrics=test_metrics, n_labels=ctx.require("n_labels"), ) model: "BaseModel" = factory.model diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 49f34add3..caf5208d1 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -5,6 +5,7 @@ from DashAI.back.core.schema_fields import ( BaseSchema, + bool_field, component_field, schema_field, string_field, @@ -65,6 +66,36 @@ class FitModelSchema(BaseSchema): zh="目标指标", ), ) # type: ignore + validation_during_fit: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether the validation partition is handed to the model while " + "fitting. Models use it to watch training and stop early. Turn it " + "off when the same partition is what the fit will be scored on.", + es="Si la partición de validación se entrega al modelo durante el " + "ajuste. Los modelos la usan para vigilar el entrenamiento y " + "detenerse antes. Desactivar cuando esa misma partición es sobre " + "la que se va a evaluar el ajuste.", + pt="Se a partição de validação é entregue ao modelo durante o " + "ajuste. Os modelos usam-na para acompanhar o treino e parar mais " + "cedo. Desative quando essa mesma partição for aquela sobre a qual " + "o ajuste será avaliado.", + de="Ob die Validierungspartition dem Modell beim Fitten übergeben " + "wird. Modelle nutzen sie, um das Training zu beobachten und früh " + "abzubrechen. Abschalten, wenn genau diese Partition den Fit " + "bewerten soll.", + zh="拟合时是否将验证分区交给模型。模型用它监控训练并提前停止。" + "当该分区正是用于评估此次拟合时，请关闭。", + ), + alias=MultilingualString( + en="Validate while fitting", + es="Validar durante el ajuste", + pt="Validar durante o ajuste", + de="Beim Fitten validieren", + zh="拟合时验证", + ), + ) # type: ignore class FitModelUnit(BaseUnit): @@ -160,23 +191,23 @@ def validate(self, ctx: ExecutionContext) -> None: self._resolve_search() def execute(self, ctx: ExecutionContext) -> None: - import os - import pickle - - from kink import di - - config = di["config"] - model = ctx.require("model") x = ctx.require("x") y = ctx.require("y") run_id = self.config["run_id"] optimizable_parameters = ctx.require("optimizable_parameters") + # The model is pointed at the data it is about to be fitted on, here + # rather than where it was built: over folds this unit runs once per + # partition, and the metric methods read these attributes off the + # instance to decide what they are scoring. + model.x_data = x + model.y_data = y + plot_paths = [] try: if not optimizable_parameters: - model.train(x["train"], y["train"], x["validation"], y["validation"]) + self._fit(model, x, y) else: # Memoized: validate() resolved these already, and resolving # again here would be the same lookup. @@ -205,9 +236,19 @@ def execute(self, ctx: ExecutionContext) -> None: factory.update_parameters(old_parameters, best_params), ) - # Generate hyperparameter plot + # Resolved here and not at the top of the method: the runs + # directory is only needed to name the plots a search produces, + # so a fit without one has no reason to require it of whatever + # is running it. + import os + import pickle + + from kink import di + from DashAI.back.core.artifacts import normalize_artifacts + config = di["config"] + trials = optimizer.get_trials_values() plot_filenames, plots = optimizer.create_plots( trials, @@ -233,6 +274,31 @@ def execute(self, ctx: ExecutionContext) -> None: ctx.put("model", model) ctx.put_ref("plot_paths", plot_paths) + def _fit(self, model, x, y) -> None: + """Fit the model on the training partition of the data it was given. + + Whether the validation partition goes with it is a policy and not a + shape: a model uses it to watch the fit and stop early, which is what + an ordinary holdout run wants, and which is exactly wrong when that + same partition is what the fit will be scored on -- a fold is scored on + the rows it held back, so handing them over would be measuring the fit + on data it was allowed to watch. + + The key may also simply not be there. The trailing entry a fold + splitter produces holds the pooled rows and the reserved ones and has + no validation partition at all, so there is nothing to hand over even + where the policy would allow it. + """ + # ``.get`` with the schema's own placeholder, the way the other units + # read a declared optional field: a caller that builds this unit by + # hand -- a job, a test -- should not have to name a policy it is happy + # to leave alone, and the ordinary answer is the ordinary holdout one. + fit_with_validation = self.config.get("validation_during_fit", True) + if fit_with_validation and "validation" in x: + model.train(x["train"], y["train"], x["validation"], y["validation"]) + else: + model.train(x["train"], y["train"]) + @staticmethod def _assert_model_keeps_its_runtime_state(model) -> None: """Fail loudly if the optimizer returned a model that cannot be scored. diff --git a/tests/back/api/test_model_job_as_a_graph.py b/tests/back/api/test_model_job_as_a_graph.py index 8e2794735..bf299b8de 100644 --- a/tests/back/api/test_model_job_as_a_graph.py +++ b/tests/back/api/test_model_job_as_a_graph.py @@ -304,15 +304,22 @@ def test_the_whole_graph_runs_to_completion(finished_pipeline_run): assert node["end_time"] is not None, node_id -def test_six_drawn_edges_expand_into_fifteen_wires(finished_pipeline_run): +def test_six_drawn_edges_expand_into_thirteen_wires(finished_pipeline_run): """The granularity problem, measured. The unit contract is finer than a canvas can draw: FitModelUnit alone requires seven keys. One drawn edge carries every key its two units agree on, which is what makes six shapes on a canvas enough for this graph. + + It was fifteen while BuildModelUnit still took the data. It stopped taking + it once a model could be fitted over folds, where the data changes on every + iteration and binding one partition at construction would leave the metrics + describing whichever fold happened to be built with. The two wires that + disappeared are ``x`` and ``y`` from prep to build; whoever fits the model + points it at the data now. """ edges = finished_pipeline_run["edges"] - assert len(edges) == 15 + assert len(edges) == 13 carried = {} for edge in edges: @@ -320,7 +327,7 @@ def test_six_drawn_edges_expand_into_fifteen_wires(finished_pipeline_run): assert carried == { ("load", "prep"): {"dataset", "dataset_id"}, - ("prep", "build"): {"x", "y", "n_labels", "task_name"}, + ("prep", "build"): {"n_labels", "task_name"}, ("prep", "fit"): {"x", "y", "task"}, ("build", "fit"): { "model", diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py index b35af3f65..391cbc268 100644 --- a/tests/back/units/test_fit_model_unit.py +++ b/tests/back/units/test_fit_model_unit.py @@ -163,3 +163,95 @@ def test_a_model_detached_from_its_data_is_refused_even_with_no_run(): def test_a_model_that_kept_its_data_passes_without_a_run(): FitModelUnit._assert_model_keeps_its_runtime_state(_Attached()) + + +# --------------------------------------------------------------------------- # +# What the fit is allowed to look at +# --------------------------------------------------------------------------- # + + +class _RecordingModel: + """Records the arguments of every fit, and nothing else.""" + + def __init__(self): + self.fits = [] + self.x_data = None + self.y_data = None + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.fits.append( + {"train": x_train, "validation": x_validation}, + ) + + +def _fit_context(model, x, y): + ctx = ExecutionContext() + ctx.put("model", model) + ctx.put("x", x) + ctx.put("y", y) + ctx.put("optimizable_parameters", []) + ctx.put("factory", object()) + ctx.put_ref("model_parameters", {}) + ctx.put("task", object()) + return ctx + + +_HOLDOUT = {"train": "x-train", "validation": "x-val", "test": "x-test"} +#: What a fold splitter's trailing entry looks like: the pooled rows and the +#: reserved ones, and no validation partition at all. +_POOLED = {"train": "x-pool", "test": "x-reserved"} + + +def test_an_ordinary_fit_hands_the_validation_partition_over(): + """Models use it to watch the fit and stop early, which holdout wants.""" + model = _RecordingModel() + + _unit()(_fit_context(model, _HOLDOUT, _HOLDOUT)) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + + +def test_a_fit_that_will_be_scored_on_validation_does_not_look_at_it(): + """A fold is scored on the rows it held back from its own training. + + Handing them to the fit would measure it on data it was allowed to watch, + and nothing about that failure raises -- the score simply comes out better + than the model deserves. + """ + model = _RecordingModel() + unit = _unit() + unit.config["validation_during_fit"] = False + + unit(_fit_context(model, _HOLDOUT, _HOLDOUT)) + + assert model.fits == [{"train": "x-train", "validation": None}] + + +def test_a_partition_set_without_a_validation_split_still_fits(): + """The trailing entry of a fold splitter has nothing to hand over. + + Reading ``x["validation"]`` unconditionally, as this unit used to, is a + KeyError on the very partition set that fits the model which gets kept. + """ + model = _RecordingModel() + + _unit()(_fit_context(model, _POOLED, _POOLED)) + + assert model.fits == [{"train": "x-pool", "validation": None}] + + +def test_the_fit_points_the_model_at_the_data_it_is_being_fitted_on(): + """Built once, fitted many times: the data comes with the fit, not the build. + + The metric methods read these attributes off the instance to decide what + they are scoring, so over folds they have to follow the iteration. + """ + model = _RecordingModel() + # Same partition names on both sides, as a splitter always produces, and + # distinguishable values so the two cannot be confused for one another. + y = {name: value.replace("x-", "y-") for name, value in _HOLDOUT.items()} + + _unit()(_fit_context(model, _HOLDOUT, y)) + + assert model.x_data is _HOLDOUT + assert model.y_data is y From 4906f9f1fe7cbcebdd34215d363f40adf91fb8eb Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 17:11:42 -0300 Subject: [PATCH 18/28] Make the search measure an objective the fitting unit builds BaseOptimizer.optimize took the task as its sixth argument and did the fitting and the scoring inline. develop replaced that argument with a callable, because cross-validation needs a trial to mean k fits rather than one, and the optimizer has no business knowing which. FitModelUnit was still passing the task into that position -- a silent mismatch, since a task is not callable, so it would have surfaced from inside a trial rather than from the call. The objective is now built by the unit that fits: one fit of the training partition and one score of the validation partition. That is what makes the same search reusable over anything that can be fitted and scored, which is the whole point of the inversion -- the fold sibling will hand it a loop instead, and nothing in the optimizer changes. The trial metrics move with it, and that is the part worth noticing. They were written by the optimizer, which meant the search decided what counted as a scored partition. It is a property of the thing being fitted: a partition with no metrics configured writes nothing, because calculate_metrics finds nothing to score and returns. So the objective writes them. `task` leaves FitModelUnit.REQUIRES, since nothing reads it there any more -- the audit would have caught it otherwise. It is still produced and still consumed, by the local explanation unit. The graph measurement drops from thirteen wires to twelve and says which one went and why. Covered directly rather than through a job: the graph test's model declares no optimizable parameters, so the search branch never runs there, and the orchestration net exercises develop's ModelJob rather than these units. Three tests pin what the objective computes, what it logs, and -- separately -- that it is what reaches the optimizer, because passing the wrong sixth argument is invisible until something calls it. 976 passed across units, dag, spike and api; the one failure is the CV-aware explanation indexes still to be ported. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/units/fit_model_unit.py | 45 ++++++- tests/back/api/test_model_job_as_a_graph.py | 11 +- tests/back/units/test_fit_model_unit.py | 126 +++++++++++++++++++- 3 files changed, 176 insertions(+), 6 deletions(-) diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index caf5208d1..778733e83 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -125,7 +125,6 @@ class FitModelUnit(BaseUnit): "model_parameters", "x", "y", - "task", ) PROVIDES = ("model", "plot_paths") RUNTIME_PARAMS = ("run_id", "artifact_prefix") @@ -220,7 +219,7 @@ def execute(self, ctx: ExecutionContext) -> None: y, optimizable_parameters, goal_metric, - ctx.require("task"), + self._score_one_trial, ) model = optimizer.get_model() best_params = optimizer.get_best_params() @@ -274,6 +273,48 @@ def execute(self, ctx: ExecutionContext) -> None: ctx.put("model", model) ctx.put_ref("plot_paths", plot_paths) + def _score_one_trial(self, model, x, y, metric) -> float: + """Fit the model once and score it, for one point of the search. + + This is what the optimizer measures. It used to be the optimizer's own + business: the objective fitted and scored inline, and the sixth + argument of ``optimize`` was the task. It is now a callable the caller + supplies, which is the change that lets the search be reused over + anything that can be fitted and scored -- a single split here, a whole + set of folds in the sibling -- without the optimizer knowing which. + + The trial metrics are written here rather than by the optimizer for the + same reason: what counts as a scored partition is a property of the + thing being fitted, not of the search. A partition with no metrics + configured for it writes nothing, because ``calculate_metrics`` finds + nothing to score and returns. + + Parameters + ---------- + model : BaseModel + The instance the optimizer has just set this trial's parameters on. + x, y : DatasetDict + The partitions this fit may use. + metric : BaseMetric + The metric class the search is optimizing, already unwrapped from + its registry entry by the optimizer. + + Returns + ------- + float + The score on the validation partition, which is the objective. + """ + from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum + + self._fit(model, x, y) + + model.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.TRIAL) + model.calculate_metrics(split=SplitEnum.VALIDATION, level=LevelEnum.TRIAL) + + predictions = model.predict(x["validation"]) + expected = model.prepare_output(y["validation"], is_fit=False) + return metric.score(expected, predictions) + def _fit(self, model, x, y) -> None: """Fit the model on the training partition of the data it was given. diff --git a/tests/back/api/test_model_job_as_a_graph.py b/tests/back/api/test_model_job_as_a_graph.py index bf299b8de..7a6240f2c 100644 --- a/tests/back/api/test_model_job_as_a_graph.py +++ b/tests/back/api/test_model_job_as_a_graph.py @@ -304,7 +304,7 @@ def test_the_whole_graph_runs_to_completion(finished_pipeline_run): assert node["end_time"] is not None, node_id -def test_six_drawn_edges_expand_into_thirteen_wires(finished_pipeline_run): +def test_six_drawn_edges_expand_into_twelve_wires(finished_pipeline_run): """The granularity problem, measured. The unit contract is finer than a canvas can draw: FitModelUnit alone @@ -317,9 +317,14 @@ def test_six_drawn_edges_expand_into_thirteen_wires(finished_pipeline_run): describing whichever fold happened to be built with. The two wires that disappeared are ``x`` and ``y`` from prep to build; whoever fits the model points it at the data now. + + The twelfth went when the hyperparameter search stopped being handed the + task. It was the sixth argument of ``optimize`` while the optimizer fitted + and scored inline; the objective is now a callable the fitting unit builds, + so the task is no longer anything the search needs to be told. """ edges = finished_pipeline_run["edges"] - assert len(edges) == 13 + assert len(edges) == 12 carried = {} for edge in edges: @@ -328,7 +333,7 @@ def test_six_drawn_edges_expand_into_thirteen_wires(finished_pipeline_run): assert carried == { ("load", "prep"): {"dataset", "dataset_id"}, ("prep", "build"): {"n_labels", "task_name"}, - ("prep", "fit"): {"x", "y", "task"}, + ("prep", "fit"): {"x", "y"}, ("build", "fit"): { "model", "factory", diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py index 391cbc268..e63d3c282 100644 --- a/tests/back/units/test_fit_model_unit.py +++ b/tests/back/units/test_fit_model_unit.py @@ -171,10 +171,11 @@ def test_a_model_that_kept_its_data_passes_without_a_run(): class _RecordingModel: - """Records the arguments of every fit, and nothing else.""" + """Records what was asked of it, and does nothing else.""" def __init__(self): self.fits = [] + self.logged = [] self.x_data = None self.y_data = None @@ -183,6 +184,15 @@ def train(self, x_train, y_train, x_validation=None, y_validation=None): {"train": x_train, "validation": x_validation}, ) + def predict(self, x_data): + return f"predictions-for-{x_data}" + + def prepare_output(self, y_data, is_fit=False): + return f"expected-from-{y_data}" + + def calculate_metrics(self, split, level, **kwargs): + self.logged.append((split, level)) + def _fit_context(model, x, y): ctx = ExecutionContext() @@ -255,3 +265,117 @@ def test_the_fit_points_the_model_at_the_data_it_is_being_fitted_on(): assert model.x_data is _HOLDOUT assert model.y_data is y + + +# --------------------------------------------------------------------------- # +# The objective the search measures +# --------------------------------------------------------------------------- # + + +class _NamedMetric: + """Scores by naming what it was given, so the arguments can be checked.""" + + @staticmethod + def score(expected, predictions): + return f"{expected}|{predictions}" + + +def test_one_trial_is_a_fit_and_a_score_of_the_validation_partition(): + """What the optimizer measures, and where the numbers come from. + + The objective used to be the optimizer's own business: it fitted and scored + inline, and the sixth argument of ``optimize`` was the task. Making it a + callable the unit supplies is what lets the same search be reused over + anything that can be fitted and scored. + """ + model = _RecordingModel() + + score = _unit()._score_one_trial(model, _HOLDOUT, _HOLDOUT, _NamedMetric) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + assert score == "expected-from-x-val|predictions-for-x-val" + + +def test_one_trial_logs_the_metrics_of_that_trial(): + """Written by the objective, not by the optimizer. + + What counts as a scored partition is a property of the thing being fitted + rather than of the search, so it is decided here. A partition with no + metrics configured writes nothing, because ``calculate_metrics`` finds + nothing to score and returns. + """ + from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum + + model = _RecordingModel() + + _unit()._score_one_trial(model, _HOLDOUT, _HOLDOUT, _NamedMetric) + + assert model.logged == [ + (SplitEnum.TRAIN, LevelEnum.TRIAL), + (SplitEnum.VALIDATION, LevelEnum.TRIAL), + ] + + +class _RecordingOptimizer: + """Stands in for a real optimizer to watch what it is handed.""" + + last_call = None + + def optimize(self, model, x, y, parameters, metric, strategy): + type(self).last_call = { + "model": model, + "parameters": parameters, + "metric": metric, + "strategy": strategy, + } + # A real optimizer leaves the model fitted at the best point it found. + strategy(model, x, y, _NamedMetric) + + def get_model(self): + return type(self).last_call["model"] + + def get_best_params(self): + return {} + + def get_trials_values(self): + return [] + + def create_plots(self, trials, run_id, n_params, goal_metric, artifact_prefix): + return [], [] + + +class _Factory: + @staticmethod + def update_parameters(old, best): + return dict(old) + + +def test_the_search_is_handed_the_units_own_objective(tmp_path): + """The wiring, pinned separately from what the objective computes. + + Passing the wrong sixth argument is silent until the optimizer calls it: + the task used to sit in that position, and a task is not callable, so the + mistake surfaced from inside a trial rather than from the call. + """ + registry = { + "RecordingOptimizer": {"class": _RecordingOptimizer}, + "Accuracy": {"class": _NamedMetric, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + # A search names the plots it produces after the run, so the runs directory + # is a real dependency of this path even when the double produces none. + di["config"] = {"RUNS_PATH": str(tmp_path)} + try: + model = _RecordingModel() + ctx = _fit_context(model, _HOLDOUT, _HOLDOUT) + ctx.put("optimizable_parameters", [("obj", "C", (0, 1), "number")]) + ctx.put("factory", _Factory) + + unit = _unit(optimizer_name="RecordingOptimizer") + unit(ctx) + + assert _RecordingOptimizer.last_call["strategy"] == unit._score_one_trial + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + finally: + del di["component_registry"] + del di["config"] From cb92af6b963075ac97d73911233ec18eda7182b9 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 17:31:42 -0300 Subject: [PATCH 19/28] Ask the splitter which partitions a run has, when explaining it ExplainerJob read train_indexes, test_indexes and val_indexes straight off Run.split_indexes. That is the shape a holdout run stores. A cross-validated one stores an entry per fold plus the pooled rows and the reserved ones, so explaining such a run raised KeyError inside the wrapper that reports a preparation failure -- the user was told the dataset could not be prepared, which is true and is not the reason. develop had already built what this needs: explainable_indexes asks the splitter that produced the run which partitions it has and what they are called, and maps whichever answer it gives onto the three slots the explainers are built from. A splitter added later needs no change here, and a fold run is explained on the rows no fold ever saw. Resolved in the job rather than in the unit. Unpacking the JSON column of a row is an artifact of how the column is stored rather than part of the transformation, and deciding which splitter wrote it is the same kind of unpacking. The unit's contract does not change: it still requires split_indexes, and still gets the three lists it always did. A payload that does not match its splitter now says so instead of reaching the generic wrapper. The old message read as a problem with the dataset rather than with the run's own record of how it was split, so the test that pinned it is updated with the reason. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/job/explainer_job.py | 36 +++++++++++++++++++++++++++- tests/back/api/test_explainer_job.py | 26 +++++++++++++++----- 2 files changed, 55 insertions(+), 7 deletions(-) diff --git a/DashAI/back/job/explainer_job.py b/DashAI/back/job/explainer_job.py index e4328c03c..adc0898df 100644 --- a/DashAI/back/job/explainer_job.py +++ b/DashAI/back/job/explainer_job.py @@ -12,6 +12,10 @@ Run, ) from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.splitters.splits_payload import ( + explainable_indexes, + splitter_class_for, +) 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 @@ -210,12 +214,42 @@ def run( # "cannot prepare" message. prepare.validate(ctx) + try: + # Which partitions a run has, and what they are called, is + # the splitter's answer and not something to read off the + # shape of the payload: a holdout run stores one flat + # mapping and a cross-validated one stores an entry per + # fold plus the pooled rows. The splitter maps whichever it + # produced onto the three slots the explainers are built + # from, so a splitter added later needs no change here. + from kink import di + + splitter_class = splitter_class_for( + json.loads(model_session.splits), di["component_registry"] + ) + train_idx, evaluation_idx, val_idx = explainable_indexes( + splitter_class, json.loads(run.split_indexes) + ) + except ValueError as e: + # Reported as itself: it says which of the two happened -- + # a payload that does not match its splitter, or a run that + # fitted on every row it had and so has nothing to explain. + log.exception(e) + raise JobError(str(e)) from e + 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)) + ctx.put_ref( + "split_indexes", + { + "train_indexes": train_idx, + "test_indexes": evaluation_idx, + "val_indexes": val_idx, + }, + ) prepare(ctx) except Exception as e: log.exception(e) diff --git a/tests/back/api/test_explainer_job.py b/tests/back/api/test_explainer_job.py index 3bb632fb3..cb2485fdd 100644 --- a/tests/back/api/test_explainer_job.py +++ b/tests/back/api/test_explainer_job.py @@ -40,6 +40,7 @@ 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.splitters.holdout import HoldoutSplitter from DashAI.back.tasks.base_task import BaseTask INPUT_COLUMNS = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"] @@ -187,6 +188,9 @@ def setup_test_registry(client, monkeypatch: pytest.MonkeyPatch): UninstantiableGlobalExplainer, DummyLocalExplainer, ExplainerJob, + # The job asks the splitter that produced the run which partitions + # it has, so the one the session names has to be resolvable. + HoldoutSplitter, ] ) @@ -573,11 +577,17 @@ def test_an_unknown_task_name_is_reported_by_name(client, run_id, model_session_ 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. +def test_incomplete_split_indexes_name_the_splitter_they_disagree_with( + client, run_id, dataset_1 +): + """A payload that does not match its splitter says so, and stops there. + + It used to reach the generic preparation wrapper, which told the user that + the dataset could not be prepared -- true, but not the reason, and it reads + as a problem with the dataset rather than with the run's own record of how + it was split. The partitions are now resolved by asking the splitter that + produced the run, before anything is prepared, so the message is the + mismatch itself. """ explainer_id = _create_global_explainer(client, run_id) @@ -586,9 +596,13 @@ def test_incomplete_split_indexes_report_a_preparation_error(client, run_id, dat db.get(Run, run_id).split_indexes = json.dumps({"train_indexes": [0, 1]}) db.commit() - with pytest.raises(JobError, match="Can not prepare dataset"): + with pytest.raises(JobError) as error: ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + assert str(error.value) == ( + "The run's split indexes do not match the splitter that produced it, " + "so there is no data to explain." + ) assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( ExplainerStatus.ERROR ) From 30a39bfc00ca225c0f1e00baf0db636279adde74 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 17:43:40 -0300 Subject: [PATCH 20/28] Replace the saved model only once the new one is complete SaveModelUnit called model.save straight at the destination. A save that died partway left a truncated artifact there, and the row went on pointing at it as if it were a model -- the failure is only visible later, when something tries to load it. develop had already fixed this in the job it kept, with atomic_save_path: the model is handed a temporary sibling path, and what it leaves there is moved into place once it returns. The temporary path is handed over rather than derived here because only the model knows whether it writes a single file or a directory of weights. Two consequences worth having in the tests. The model no longer sees the final path, so what is asserted is where the artifact ended up rather than what the model was told -- which is what the caller and the row care about anyway. And a double that recorded a path without writing anything now fails, correctly: leaving nothing to move is the same thing a model that silently saved nothing would do, and it should be reported rather than hidden. 978 passed across units, dag, spike and api. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/units/save_model_unit.py | 9 +++- tests/back/units/test_save_model_unit.py | 52 ++++++++++++++++++++++-- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/DashAI/back/units/save_model_unit.py b/DashAI/back/units/save_model_unit.py index fa808b61f..97ce74c6a 100644 --- a/DashAI/back/units/save_model_unit.py +++ b/DashAI/back/units/save_model_unit.py @@ -3,6 +3,7 @@ import logging import re +from DashAI.back.core.atomic import atomic_save_path from DashAI.back.job.base_job import JobError from DashAI.back.units.base_unit import BaseUnit from DashAI.back.units.context import ExecutionContext @@ -57,7 +58,13 @@ def execute(self, ctx: ExecutionContext) -> None: model_path = os.path.join( config["RUNS_PATH"], self.config["artifact_prefix"] ) - model.save(model_path) + # Written aside and moved into place, so a save that dies halfway + # leaves the previous artifact intact instead of a truncated one + # that the row still points at. The temporary path is handed to the + # model rather than derived here because only the model knows + # whether it writes a file or a directory of weights. + with atomic_save_path(model_path) as tmp_path: + model.save(str(tmp_path)) except Exception as e: log.exception(e) raise JobError( diff --git a/tests/back/units/test_save_model_unit.py b/tests/back/units/test_save_model_unit.py index 7aeb1558c..ff98c22a0 100644 --- a/tests/back/units/test_save_model_unit.py +++ b/tests/back/units/test_save_model_unit.py @@ -11,13 +11,22 @@ class _RecordingModel: - """Stands in for a trained model, recording where it was asked to go.""" + """Stands in for a trained model, recording where it was asked to go. + + It writes a real file at that path, because the save is now made atomic: + the model is handed a temporary path and what it leaves there is moved + into place. A double that recorded without writing would leave nothing to + move, which is the same failure a model that silently saved nothing would + produce -- and is exactly what should be reported rather than hidden. + """ def __init__(self) -> None: self.saved_to = None def save(self, path) -> None: self.saved_to = path + with open(path, "w", encoding="utf-8") as file: + file.write("a saved model") @pytest.fixture(name="runs_path") @@ -34,8 +43,13 @@ def test_the_model_lands_in_a_directory_named_by_the_prefix(runs_path): SaveModelUnit(artifact_prefix="pipeline-3-save")(ctx) - assert model.saved_to == os.path.join(runs_path, "pipeline-3-save") - assert ctx.require("model_path") == model.saved_to + # Where the artifact ended up, not where the model was told to write. The + # save is atomic, so the model writes to a temporary sibling and what it + # left there is moved into place; the destination is what the caller and + # the row care about. + expected = os.path.join(runs_path, "pipeline-3-save") + assert os.path.exists(expected) + assert ctx.require("model_path") == expected def test_a_numeric_run_id_is_a_valid_prefix(runs_path): @@ -51,7 +65,8 @@ def test_a_numeric_run_id_is_a_valid_prefix(runs_path): SaveModelUnit(artifact_prefix="17")(ctx) - assert model.saved_to == os.path.join(runs_path, "17") + assert os.path.exists(os.path.join(runs_path, "17")) + assert ctx.require("model_path") == os.path.join(runs_path, "17") @pytest.mark.parametrize( @@ -89,3 +104,32 @@ def test_the_unit_still_needs_a_model(runs_path): """``run_id`` left REQUIRES; ``model`` did not.""" with pytest.raises(UnitContractError, match="'model'"): SaveModelUnit(artifact_prefix="1")(ExecutionContext()) + + +def test_a_save_that_dies_halfway_leaves_the_previous_artifact_alone(runs_path): + """The reason the save is atomic. + + A model that raises partway through writing used to leave a truncated + artifact at the destination, which the row then pointed at as if it were a + model. Now the destination is only replaced once the new one is complete. + """ + destination = os.path.join(runs_path, "42") + with open(destination, "w", encoding="utf-8") as file: + file.write("the previous model") + + class _DyingModel: + def save(self, path): + with open(path, "w", encoding="utf-8") as file: + file.write("half a model") + raise RuntimeError("out of disk") + + ctx = ExecutionContext() + ctx.put("model", _DyingModel()) + + with pytest.raises(JobError, match="Model saving failed"): + SaveModelUnit(artifact_prefix="42")(ctx) + + with open(destination, encoding="utf-8") as file: + assert file.read() == "the previous model" + # And nothing half-written was left lying next to it. + assert os.listdir(runs_path) == ["42"] From 7c708d47f19822cf54d277ce7c1620d998817b99 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 17:54:35 -0300 Subject: [PATCH 21/28] Prepare the data through the units, on both training paths ModelJob had its own copy of everything before the model is fitted: loading the dataset, resolving the task, validating the dataset against it, counting the labels, separating features from targets, resolving the metrics and the model class, checking the downloads, and calling the splitter. The units had the same steps. That was most of the duplication this reconciliation exists to remove, and it goes in one piece rather than one path at a time -- holdout and cross-validation differ only in which unit prepares the data. _prepare_dataset_and_components now does what only it can: read the rows, unpack the JSON columns stored on them, and choose which unit prepares the data. That choice follows from how the splitter carves the dataset, which the splitter declares -- it is a choice of unit rather than a flag on one, because the two publish different shapes. The file loses eighty-four lines. The evaluation strategy is built in run() now rather than in the helper, because it takes the factory the build unit produced. Two orderings are deliberate and were not obvious: BuildModelUnit.validate runs before the data is partitioned, and the unit itself after. The download gate lives in validate, and a model that cannot be trained should be reported as that rather than surfacing later as a splitting failure -- the same reason ModelJob has always resolved the optimizer before changing the run's status. The splitter class is resolved in the helper and again inside the unit. The helper needs it to know which unit to build, and the message a missing one produces belongs to the preparation step. The unit resolves its own because it must work for a caller that is not this job. Both regression nets pass unchanged -- 46 tests, including the verbatim text of every error branch, which is what says the messages did not drift. 978 across units, dag, spike and api. The evaluation strategy still owns the training loop. That is the next piece. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/job/model_job.py | 282 ++++++++++++----------------------- 1 file changed, 99 insertions(+), 183 deletions(-) diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 96e00de9f..ee4b3cafb 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -1,25 +1,24 @@ import logging -from typing import TYPE_CHECKING, Any, Dict, List +from typing import TYPE_CHECKING, Any, Dict from kink import inject from sqlalchemy import exc from DashAI.back.core.atomic import atomic_save_path from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run -from DashAI.back.dependencies.downloads.nested import missing_downloads from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.metrics.base_metric import BaseMetric -from DashAI.back.models.model_factory import ModelFactory from DashAI.back.optimizers.base_optimizer import BaseOptimizer -from DashAI.back.splitters.base_splitter import BaseSplitter from DashAI.back.splitters.splits_payload import normalize_splits_payload -from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_model_unit import BuildModelUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit +from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker - from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -105,12 +104,12 @@ def run( # Get the necessary parameters run_id: int = self.kwargs["run_id"] + ctx = ExecutionContext() + with session_factory() as db: run: Run = db.get(Run, run_id) # Without this the next line raises AttributeError on None, which - # reaches the user as a stack trace rather than as the reason. The - # guard came from the atomized job and is kept here until this - # method is rebuilt on the units. + # reaches the user as a stack trace rather than as the reason. if not run: raise JobError(f"Run {run_id} does not exist in DB.") run.huey_id = self.kwargs.get("huey_id", None) @@ -118,10 +117,31 @@ def run( self.report_progress(0.05, "Preparing data") try: try: - # Get the dataset and components prepared for the model training + # What is left in the helper is reading the configuration + # this run was created with off its rows. The work that + # configuration describes is done by the units below. preparation_results = self._prepare_dataset_and_components( run_id=run_id, db=db, component_registry=component_registry ) + model_session: ModelSession = preparation_results["model_session"] + prepare = preparation_results["prepare_unit"] + + LoadDatasetUnit(dataset_id=model_session.dataset_id)(ctx) + + build_model = BuildModelUnit( + model={ + "component": run.model_name, + "params": run.parameters, + }, + train_metrics=model_session.train_metrics, + validation_metrics=model_session.validation_metrics, + test_metrics=model_session.test_metrics, + run_id=run_id, + ) + # The download gate lives in validate(), and running it here + # keeps a model that cannot be trained from being reported + # as a splitting failure further down. + build_model.validate(ctx) except Exception as e: log.exception(e) raise JobError( @@ -129,23 +149,32 @@ def run( ) from e try: - # Get splits from the splitter - splitter: BaseSplitter = preparation_results["splitter"] - # Get the dataset splits between input columns and output column - X, Y = preparation_results["X"], preparation_results["Y"] + # Preparing the dataset for the task and partitioning it are + # one step: how many partitions there are and what they are + # called is the splitter's answer, and the pair of units + # differ only in which family of splitters they offer and in + # the shape they publish for it. + prepare(ctx) - # Get x,y but now splitted with train, validation and test indexes - # each one, and the indexes used for the splits - x, y, splits = splitter.split(X, Y) + x = ctx.get("x") if ctx.has("x") else ctx.require("x_folds") + y = ctx.get("y") if ctx.has("y") else ctx.require("y_folds") # save the obtained splits into the database - run.split_indexes = json.dumps(splits) + run.split_indexes = json.dumps(ctx.require("split_indexes")) except Exception as e: log.exception(e) raise JobError( f"Error splitting the dataset for run {run_id}: {e}", ) from e + try: + build_model(ctx) + except Exception as e: + log.exception(e) + raise JobError( + f"Error preparing dataset and components for run {run_id}: {e}", + ) from e + try: run.set_status_as_started() db.commit() @@ -160,9 +189,14 @@ def run( # Hyperparameter Tunning plot_paths = [] - evaluation_estrategy: BaseEvaluationStrategy = preparation_results[ - "evaluation_strategy" - ] + # Built here rather than in the helper because it takes + # the factory the build unit produced. + strategy_class = preparation_results["evaluation_strategy_class"] + evaluation_estrategy: BaseEvaluationStrategy = strategy_class( + factory=ctx.require("factory"), + optimizer=preparation_results["optimizer"], + goal_metric=preparation_results["goal_metric"], + ) evaluation_estrategy.set_progress_reporter(self.report_progress) model, plot_paths = evaluation_estrategy.execute( @@ -227,53 +261,48 @@ def run( db.commit() raise e finally: + ctx.clear_cache() gc.collect() def _prepare_dataset_and_components( self, run_id: int, db, component_registry ) -> Dict[str, Any]: - """Prepare the dataset, task, splitter, metrics, model, and evaluation strategy. + """Read the configuration this run was created with, off its rows. - This helper resolves the persisted training configuration for a run, - loads the associated dataset from disk, prepares it for the selected - task, instantiates the required components from the component registry, - and builds the model factory together with the evaluation strategy. + What is resolved here is what the units cannot: rows, the JSON columns + stored on them, and the choice of which unit prepares the data. The + work those rows describe -- loading the dataset, validating it against + the task, separating features from targets, partitioning them and + building the model -- belongs to the units and happens in ``run``. Parameters ---------- run_id : int - Identifier of the training run whose configuration and artifacts must - be loaded. + Identifier of the training run whose configuration must be read. db : object - Database access object used to retrieve the run, model session, and - related persisted entities. + Database session used to retrieve the run and its model session. component_registry : object - Registry containing the available task, splitter, metric, model, - optimizer, and evaluation strategy implementations. + Registry used to resolve the splitter, the optimizer and the + evaluation strategy the session names. Returns ------- dict - A dictionary containing the prepared input and output datasets, the - instantiated splitter, and the evaluation strategy. + The model session, the unit that will prepare and partition the + dataset, the evaluation strategy class, and the optimizer and goal + metric it will be built with. Raises ------ JobError - If the run, model session, dataset, task, splitter, metrics, model, - optimizer, or evaluation strategy cannot be resolved or instantiated. + If the run's session, its splits payload, its splitter, its + optimizer or its evaluation strategy cannot be resolved. """ import json - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - select_columns, - ) - run: Run = db.get(Run, run_id) - # Get the model session and dataset from the database model_session: ModelSession = db.get(ModelSession, run.model_session_id) if not model_session: raise JobError( @@ -285,71 +314,13 @@ def _prepare_dataset_and_components( raise JobError(f"Dataset {model_session.dataset_id} does not exist in DB.") try: - # Load dataset from the file path - loaded_dataset: "DashAIDataset" = 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: - # Get task from model session - 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: - # Prepare dataset for the task and get number of labels of the task - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=model_session.input_columns, - output_columns=model_session.output_columns, - ) - n_labels = task.num_labels( - prepared_dataset, model_session.output_columns[0] - ) - except Exception as e: - log.exception(e) - raise JobError( - f"""Can not prepare Dataset {dataset.id} - for Task {model_session.task_name}""", - ) from e - - try: - # Divide the dataset into two datasets: - # one with the input columns and another with the output column. - # This reads the prepared dataset rather than the loaded one: a - # task may reorder or otherwise adjust the rows, and forecasting - # does, sorting them by date so the temporal splitter carves real - # periods of time. Selecting from the loaded dataset would drop - # that work on the floor. - X, Y = select_columns( - prepared_dataset, - model_session.input_columns, - model_session.output_columns, - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error selecting input and output columns from dataset {dataset.id}", - ) from e - - try: - # Get splits data from model session + # Unpacking the JSON column is an artifact of how the row stores + # it, not part of the split. Sessions created before the payload + # followed the splitter schema use different keys for the seed and + # for manual indexes. splits_data = json.loads(model_session.splits) if run.split_indexes: splits_data["splitted_indexes"] = json.loads(run.split_indexes) - # Sessions created before the splits payload followed the splitter - # schema use different keys for the seed and for manual indexes. splits_data = normalize_splits_payload(splits_data) except Exception as e: log.exception(e) @@ -358,61 +329,28 @@ def _prepare_dataset_and_components( ) from e try: - # Get the splitter class from the registry and split the dataset splitter_name = splits_data.get("splitter_name", None) - splitter: BaseSplitter = component_registry[splitter_name]["class"]( - splits_data=splits_data, - ) + splitter_class = component_registry[splitter_name]["class"] except Exception as e: log.exception(e) raise JobError( - f"""Unable to find Splitter with name - {splitter_name} in registry.""", + f"Unable to find Splitter with name {splitter_name} in registry.", ) from e - try: - # Get metrics from model session - train_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.train_metrics - ] - validation_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.validation_metrics - ] - test_metrics: List[BaseMetric] = [ - component_registry[m]["class"] for m in model_session.test_metrics - ] - except Exception as e: - log.exception(e) - raise JobError( - "Unable to find metrics associated with" - f"Task {model_session.task_name} in registry", - ) from e - - try: - # Get the model class from the registry - 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 - - # Make sure the model (and any nested components) are downloaded - # before attempting to train, otherwise fail fast with a clear error. - if getattr(run_model_class, "REQUIRES_DOWNLOAD", False) and not ( - run_model_class.is_downloaded() - ): - raise JobError( - f"Model {run.model_name} is not downloaded. " - "Download it before training." - ) - nested_missing = missing_downloads(run.parameters, component_registry) - if nested_missing: - names = ", ".join(m["name"] for m in nested_missing) - raise JobError( - "These components are not downloaded. " - f"Download them before training: {names}." - ) + # Which unit prepares the data follows from how the splitter carves it, + # which the splitter declares. The two units publish different shapes, + # so this is a choice of unit and not a flag on one. + prepare_class = ( + PrepareAndFoldUnit + if getattr(splitter_class, "PARTITIONING", "holdout") == "folds" + else PrepareAndSplitUnit + ) + prepare_unit = prepare_class( + task_name=model_session.task_name, + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + splitter={"component": splitter_name, "params": splits_data}, + ) try: # Get the optimizer if defined @@ -432,32 +370,9 @@ def _prepare_dataset_and_components( ) from e try: - # Instantiate the model using the ModelFactory - # and get the optimizable parameters - factory = ModelFactory( - model=run_model_class, - params=run.parameters, - run_id=run_id, - train_metrics=train_metrics, - validation_metrics=validation_metrics, - test_metrics=test_metrics, - n_labels=n_labels, - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to instantiate model using run {run_id}", - ) from e - - try: - # Get the evaluation strategy for the model session - evaluation_strategy: BaseEvaluationStrategy = component_registry[ + evaluation_strategy_class = component_registry[ model_session.evaluation_strategy - ]["class"]( - factory=factory, - optimizer=optimizer, - goal_metric=goal_metric, - ) + ]["class"] except Exception as e: log.exception(e) raise JobError( @@ -467,8 +382,9 @@ def _prepare_dataset_and_components( ) from e return { - "X": X, - "Y": Y, - "splitter": splitter, - "evaluation_strategy": evaluation_strategy, + "model_session": model_session, + "prepare_unit": prepare_unit, + "evaluation_strategy_class": evaluation_strategy_class, + "optimizer": optimizer, + "goal_metric": goal_metric, } From c761a003fc88553c88cea7338a2977a0ed618e0c Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 18:47:43 -0300 Subject: [PATCH 22/28] Train a holdout run through the units, not the strategy HoldoutEvaluationStrategy.execute was FitModelUnit, EvaluateModelUnit and SaveModelUnit in sequence, written a second time. ModelJob now composes those three for a holdout run, and the strategy is no longer called for one. Fold runs still train through it; their loop is the next piece. Saving moved for both paths at once. The strategy hands the model back rather than leaving it in the context, so the job puts it there and one SaveModelUnit serves whichever path produced it -- which also gets the fold path the atomic replacement it did not have separately. FitModelUnit gained trial_splits, which is the SCORED_SPLITS question answered for the search. Which partitions a run records a score for is declared by the strategy the session chose, and it is not the same as which ones have metrics configured: a forecaster has training metrics and still must not be judged on the dates it was fitted on, because an in-sample fit statistic is not comparable with a forecast. The job reads that declaration and passes it on, so the strategy classes keep deciding it while the units do the work. The test partition is deliberately not an option in that field. Scoring it once per trial would let the search see it, and a model chosen with the test set in view has no honest score left to report -- so a trial may score the partition it fitted on and the one it is measured against, and nothing else. It was a default before, which is a weaker statement than a value that cannot be chosen. The runs directory leaves the job: naming the artifact is the saving unit's, and the job had been resolving it only to build a path the unit builds itself. Both nets pass unchanged, 46 tests. 992 across units, dag, spike, api and evaluation -- the strategies' own tests included, since the classes still stand and are still what declares the scored partitions. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/job/model_job.py | 91 +++++++++++++++++++---------- DashAI/back/units/fit_model_unit.py | 51 ++++++++++++++-- 2 files changed, 107 insertions(+), 35 deletions(-) diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index ee4b3cafb..6475f4b08 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -3,8 +3,8 @@ from kink import inject from sqlalchemy import exc +from sqlalchemy.orm.attributes import flag_modified -from DashAI.back.core.atomic import atomic_save_path from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy from DashAI.back.job.base_job import BaseJob, JobError @@ -12,9 +12,12 @@ from DashAI.back.splitters.splits_payload import normalize_splits_payload from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.load_dataset_unit import LoadDatasetUnit from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.save_model_unit import SaveModelUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -93,13 +96,11 @@ def run( ) -> None: import gc import json - import os from kink import di component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] # Get the necessary parameters run_id: int = self.kwargs["run_id"] @@ -184,27 +185,65 @@ def run( "Connection with the database failed", ) from e + strategy_class = preparation_results["evaluation_strategy_class"] + # Which partitions a run records a score for is declared by the + # strategy the session chose: a forecaster is not judged on the + # dates it was fitted on, so its training partition is not in + # this list even though metrics are configured for it. + scored_splits = [split.name for split in strategy_class.SCORED_SPLITS] + self.report_progress(0.2, "Training") try: - # Hyperparameter Tunning plot_paths = [] - # Built here rather than in the helper because it takes - # the factory the build unit produced. - strategy_class = preparation_results["evaluation_strategy_class"] - evaluation_estrategy: BaseEvaluationStrategy = strategy_class( - factory=ctx.require("factory"), - optimizer=preparation_results["optimizer"], - goal_metric=preparation_results["goal_metric"], - ) - - evaluation_estrategy.set_progress_reporter(self.report_progress) - model, plot_paths = evaluation_estrategy.execute( - x=x, - y=y, - run=run, - db=db, - ) + if getattr(strategy_class, "KIND", "holdout") == "holdout": + fit_model = FitModelUnit( + optimizer={ + "component": run.optimizer_name, + "params": run.optimizer_parameters, + }, + goal_metric=run.goal_metric, + run_id=run_id, + # The run names its own artifacts, which is what + # keeps the plot filenames of two runs apart inside + # the runs directory. + artifact_prefix=str(run_id), + # A trial never scores the test partition, so what + # it may record is whatever else the strategy scores. + trial_splits=[ + name for name in scored_splits if name != "TEST" + ], + ) + fit_model(ctx) + + plot_paths = ctx.require("plot_paths") + if ctx.has("best_parameters"): + run.parameters = ctx.get("best_parameters") + flag_modified(run, "parameters") + db.commit() + + self.report_progress(0.85, "Computing metrics") + EvaluateModelUnit(run_id=run_id, splits=scored_splits)(ctx) + else: + # Fold runs still train through the strategy. Their loop + # is the next piece to move; everything before and after + # it is already the units'. + evaluation_estrategy: BaseEvaluationStrategy = strategy_class( + factory=ctx.require("factory"), + optimizer=preparation_results["optimizer"], + goal_metric=preparation_results["goal_metric"], + ) + evaluation_estrategy.set_progress_reporter(self.report_progress) + model, plot_paths = evaluation_estrategy.execute( + x=x, + y=y, + run=run, + db=db, + ) + # The strategy hands the model back rather than leaving + # it in the context, so the saving unit below serves + # both paths. + ctx.put("model", model) except Exception as e: log.exception(e) raise JobError( @@ -227,18 +266,10 @@ def run( ) from e self.report_progress(0.95, "Saving model") - try: - run_path = os.path.join(config["RUNS_PATH"], str(run.id)) - with atomic_save_path(run_path) as tmp_run_path: - model.save(str(tmp_run_path)) - except Exception as e: - log.exception(e) - raise JobError( - "Model saving failed", - ) from e + SaveModelUnit(artifact_prefix=str(run_id))(ctx) try: - run.run_path = run_path + run.run_path = ctx.require("model_path") db.commit() except exc.SQLAlchemyError as e: log.exception(e) diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 778733e83..0fa60d8ce 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -3,10 +3,13 @@ import logging from typing import TYPE_CHECKING +from DashAI.back.core.enums.metrics import SplitEnum from DashAI.back.core.schema_fields import ( BaseSchema, bool_field, component_field, + enum_field, + list_field, schema_field, string_field, ) @@ -20,6 +23,12 @@ log = logging.getLogger(__name__) +#: A trial may score the partition it fitted on and the one it is measured +#: against, and nothing else. The test partition is deliberately absent: +#: scoring it once per trial would let the search see it, and a model chosen +#: with the test set in view has no honest score left to report. +TRIAL_SPLITS = ["TRAIN", "VALIDATION"] + class FitModelSchema(BaseSchema): optimizer: schema_field( @@ -66,6 +75,34 @@ class FitModelSchema(BaseSchema): zh="目标指标", ), ) # type: ignore + trial_splits: schema_field( + list_field(enum_field(enum=TRIAL_SPLITS)), + placeholder=TRIAL_SPLITS, + description=MultilingualString( + en="Partitions each trial of the search records a score for. A " + "partition the model is not meant to be judged on belongs out of " + "this list even when metrics are configured for it.", + es="Particiones para las que cada intento de la búsqueda registra " + "un puntaje. Una partición sobre la que el modelo no debe juzgarse " + "no va en esta lista aunque tenga métricas configuradas.", + pt="Partições para as quais cada tentativa da procura regista uma " + "pontuação. Uma partição sobre a qual o modelo não deve ser julgado " + "fica fora desta lista mesmo que tenha métricas configuradas.", + de="Partitionen, für die jeder Versuch der Suche einen Wert " + "festhält. Eine Partition, nach der das Modell nicht beurteilt " + "werden soll, gehört nicht in diese Liste, auch wenn Metriken für " + "sie konfiguriert sind.", + zh="搜索的每次试验为其记录分数的分区。不应据以评判模型的分区不列入此处，" + "即使已为其配置了指标。", + ), + alias=MultilingualString( + en="Trial splits", + es="Particiones por intento", + pt="Partições por tentativa", + de="Versuchspartitionen", + zh="试验分区", + ), + ) # type: ignore validation_during_fit: schema_field( bool_field(), placeholder=True, @@ -285,8 +322,12 @@ def _score_one_trial(self, model, x, y, metric) -> float: The trial metrics are written here rather than by the optimizer for the same reason: what counts as a scored partition is a property of the - thing being fitted, not of the search. A partition with no metrics - configured for it writes nothing, because ``calculate_metrics`` finds + thing being fitted, not of the search. Which ones those are is + configured, because it is not always the same as which ones have + metrics: a forecaster has training metrics and still must not be judged + on the dates it was fitted on, since an in-sample fit statistic is not + comparable with a forecast. A partition left in the list but with no + metrics configured writes nothing anyway -- ``calculate_metrics`` finds nothing to score and returns. Parameters @@ -304,12 +345,12 @@ def _score_one_trial(self, model, x, y, metric) -> float: float The score on the validation partition, which is the objective. """ - from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum + from DashAI.back.core.enums.metrics import LevelEnum self._fit(model, x, y) - model.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.TRIAL) - model.calculate_metrics(split=SplitEnum.VALIDATION, level=LevelEnum.TRIAL) + for name in self.config.get("trial_splits", TRIAL_SPLITS): + model.calculate_metrics(split=SplitEnum[name], level=LevelEnum.TRIAL) predictions = model.predict(x["validation"]) expected = model.prepare_output(y["validation"], is_fit=False) From 9c6ac25ba75d1fd11ab48ad04ddb5e5c9746ceb4 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 18:57:50 -0300 Subject: [PATCH 23/28] Extract what surrounds a fit, before there are two units doing it The cross-validation sibling needs the same optimizer resolution, the same search, the same check that the optimizer gave back the model it was handed, the same recording of the best parameters and the same plot writing. Copying that is how the two implementations this branch is removing came to exist, so it moves to fit_scope.py first and the sibling is written against it. The contract audit caught a real mistake in the first attempt, and it is worth recording because the rule reads like tidiness until it bites. The helper took the context and did its own require and put. The audit parses each unit's own source, so moving the reads out of the unit made four declared keys look unread: FitModelUnit was suddenly requiring `factory` and `model_parameters` and declaring `run_id` and `artifact_prefix` while appearing to touch none of them. A caller reading only the declarations -- the graph validator among them -- would have been told the truth by the declarations and contradicted by the audit, or worse, the declarations would have been trimmed to match. So the helper takes and returns plain values and never touches the context. Every require and put stays in the unit, and so does reading the runtime parameters. It is the mirror of the rule already written down for a helper that publishes: a ctx.put hidden in one makes a broken PROVIDES pass. 992 passed across units, dag, spike, api and evaluation. No behaviour changed: this is the same fit, moved. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/units/fit_model_unit.py | 333 ++++---------------------- DashAI/back/units/fit_scope.py | 354 ++++++++++++++++++++++++++++ 2 files changed, 399 insertions(+), 288 deletions(-) create mode 100644 DashAI/back/units/fit_scope.py diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 0fa60d8ce..018c2e0da 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -1,141 +1,32 @@ -"""Unit that fits a model, optionally searching for its hyperparameters.""" +"""Unit that fits a model on one set of partitions, optionally searching.""" import logging -from typing import TYPE_CHECKING -from DashAI.back.core.enums.metrics import SplitEnum -from DashAI.back.core.schema_fields import ( - BaseSchema, - bool_field, - component_field, - enum_field, - list_field, - schema_field, - string_field, -) -from DashAI.back.core.utils import MultilingualString +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema 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.optimizers.base_optimizer import BaseOptimizer +from DashAI.back.units.fit_scope import ( + TRIAL_SPLITS, + ModelFitScopeMixin, + goal_metric_field, + optimizer_field, + trial_splits_field, + validation_during_fit_field, +) log = logging.getLogger(__name__) -#: A trial may score the partition it fitted on and the one it is measured -#: against, and nothing else. The test partition is deliberately absent: -#: scoring it once per trial would let the search see it, and a model chosen -#: with the test set in view has no honest score left to report. -TRIAL_SPLITS = ["TRAIN", "VALIDATION"] - class FitModelSchema(BaseSchema): - optimizer: schema_field( - component_field(parent="BaseOptimizer"), - placeholder={"component": "OptunaOptimizer", "params": {}}, - description=MultilingualString( - en="Optimizer used to search for hyperparameters, along with its own " - "configuration. Only used when the model declares optimizable " - "parameters.", - es="Optimizador usado para buscar hiperparámetros, junto con su propia " - "configuración. Solo se usa cuando el modelo declara parámetros " - "optimizables.", - pt="Otimizador usado para procurar hiperparâmetros, junto com a sua " - "própria configuração. Só é usado quando o modelo declara parâmetros " - "otimizáveis.", - de="Optimierer für die Hyperparametersuche samt seiner eigenen " - "Konfiguration. Wird nur verwendet, wenn das Modell optimierbare " - "Parameter deklariert.", - zh="用于搜索超参数的优化器及其自身配置。仅当模型声明了可优化参数时使用。", - ), - alias=MultilingualString( - en="Optimizer", - es="Optimizador", - pt="Otimizador", - de="Optimierer", - zh="优化器", - ), - ) # type: ignore - goal_metric: schema_field( - string_field(), - placeholder="Accuracy", - description=MultilingualString( - en="Metric the hyperparameter search optimizes.", - es="Métrica que optimiza la búsqueda de hiperparámetros.", - pt="Métrica que a procura de hiperparâmetros otimiza.", - de="Metrik, die die Hyperparametersuche optimiert.", - zh="超参数搜索所优化的指标。", - ), - alias=MultilingualString( - en="Goal metric", - es="Métrica objetivo", - pt="Métrica objetivo", - de="Zielmetrik", - zh="目标指标", - ), - ) # type: ignore - trial_splits: schema_field( - list_field(enum_field(enum=TRIAL_SPLITS)), - placeholder=TRIAL_SPLITS, - description=MultilingualString( - en="Partitions each trial of the search records a score for. A " - "partition the model is not meant to be judged on belongs out of " - "this list even when metrics are configured for it.", - es="Particiones para las que cada intento de la búsqueda registra " - "un puntaje. Una partición sobre la que el modelo no debe juzgarse " - "no va en esta lista aunque tenga métricas configuradas.", - pt="Partições para as quais cada tentativa da procura regista uma " - "pontuação. Uma partição sobre a qual o modelo não deve ser julgado " - "fica fora desta lista mesmo que tenha métricas configuradas.", - de="Partitionen, für die jeder Versuch der Suche einen Wert " - "festhält. Eine Partition, nach der das Modell nicht beurteilt " - "werden soll, gehört nicht in diese Liste, auch wenn Metriken für " - "sie konfiguriert sind.", - zh="搜索的每次试验为其记录分数的分区。不应据以评判模型的分区不列入此处，" - "即使已为其配置了指标。", - ), - alias=MultilingualString( - en="Trial splits", - es="Particiones por intento", - pt="Partições por tentativa", - de="Versuchspartitionen", - zh="试验分区", - ), - ) # type: ignore - validation_during_fit: schema_field( - bool_field(), - placeholder=True, - description=MultilingualString( - en="Whether the validation partition is handed to the model while " - "fitting. Models use it to watch training and stop early. Turn it " - "off when the same partition is what the fit will be scored on.", - es="Si la partición de validación se entrega al modelo durante el " - "ajuste. Los modelos la usan para vigilar el entrenamiento y " - "detenerse antes. Desactivar cuando esa misma partición es sobre " - "la que se va a evaluar el ajuste.", - pt="Se a partição de validação é entregue ao modelo durante o " - "ajuste. Os modelos usam-na para acompanhar o treino e parar mais " - "cedo. Desative quando essa mesma partição for aquela sobre a qual " - "o ajuste será avaliado.", - de="Ob die Validierungspartition dem Modell beim Fitten übergeben " - "wird. Modelle nutzen sie, um das Training zu beobachten und früh " - "abzubrechen. Abschalten, wenn genau diese Partition den Fit " - "bewerten soll.", - zh="拟合时是否将验证分区交给模型。模型用它监控训练并提前停止。" - "当该分区正是用于评估此次拟合时，请关闭。", - ), - alias=MultilingualString( - en="Validate while fitting", - es="Validar durante el ajuste", - pt="Validar durante o ajuste", - de="Beim Fitten validieren", - zh="拟合时验证", - ), - ) # type: ignore + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + trial_splits: trial_splits_field() # type: ignore + validation_during_fit: validation_during_fit_field() # type: ignore -class FitModelUnit(BaseUnit): +class FitModelUnit(BaseUnit, ModelFitScopeMixin): """Train a model, running a hyperparameter search when there is one to run. Hyperparameter optimization is a fitting strategy rather than a separate @@ -145,9 +36,11 @@ class FitModelUnit(BaseUnit): ``validate`` resolves the optimizer and the goal metric so an impossible configuration is rejected before the job reports that training started. - The optimizer is configured as a component field, so its value is - ``{"component": , "params": {...}}`` and the front renders the - chosen optimizer's own form underneath. + The sibling for cross-validation is ``FitModelOverFoldsUnit``: it takes a + list of partition sets rather than one, which is a different ``REQUIRES`` + and therefore a different unit. What surrounds the fit is shared between + them; only the objective the search measures, and what happens once the + search is over, is written here. """ SCHEMA = FitModelSchema @@ -171,136 +64,47 @@ def __init__(self, **config) -> None: self._optimizer = None self._goal_metric = None - def _resolve_search(self): - """Resolve the optimizer and the goal metric, memoized on this unit. - - Kept on the instance rather than in the context on purpose. These are - this unit's own state, not something it hands to another unit: two - ``FitModelUnit`` instances sharing a context — a DAG with two training - nodes — would otherwise overwrite each other's optimizer, and the - second one would silently run the first one's. - """ - if self._optimizer is not None: - return self._optimizer, self._goal_metric - - from kink import di - - component_registry = di["component_registry"] - goal_metric_name: str = self.config["goal_metric"] - optimizer_name: str = self.config["optimizer"]["component"] - - try: - # The whole registry entry, not the class: the optimizer reads - # metadata["maximize"] from it to pick a direction. - goal_metric = component_registry[goal_metric_name] - except Exception as e: - log.exception(e) - raise JobError( - f"Metric is not compatible with the Task. {e}", - ) from e - - try: - optimizer_class = component_registry[optimizer_name]["class"] - optimizer: "BaseOptimizer" = optimizer_class( - **self.config["optimizer"]["params"] - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Error instantiating optimizer {optimizer_name}, {e}", - ) from e - - self._goal_metric = goal_metric - self._optimizer = optimizer - return optimizer, goal_metric - def validate(self, ctx: ExecutionContext) -> None: - # ctx.require, not ctx.get: "optimizable_parameters" is one of this - # unit's REQUIRES, so its absence means BuildModelUnit hasn't run yet - # — a call-order mistake, not "there is nothing to optimize". Only an - # empty value (the key present, genuinely no optimizable parameters) - # skips the optimizer/goal-metric checks below, so no registry lookup - # is needed either. - if not ctx.require("optimizable_parameters"): - return - - self._resolve_search() + # ctx.require, not ctx.get: an absent key means BuildModelUnit has not + # run yet, which is a call-order mistake rather than a model with + # nothing to optimize. + self._validate_search(ctx.require("optimizable_parameters")) def execute(self, ctx: ExecutionContext) -> None: model = ctx.require("model") x = ctx.require("x") y = ctx.require("y") - run_id = self.config["run_id"] optimizable_parameters = ctx.require("optimizable_parameters") # The model is pointed at the data it is about to be fitted on, here - # rather than where it was built: over folds this unit runs once per - # partition, and the metric methods read these attributes off the - # instance to decide what they are scoring. + # rather than where it was built: the metric methods read these + # attributes off the instance to decide what they are scoring. model.x_data = x model.y_data = y plot_paths = [] try: if not optimizable_parameters: - self._fit(model, x, y) + self._fit_kept_model(model, x, y) else: - # Memoized: validate() resolved these already, and resolving - # again here would be the same lookup. - optimizer, goal_metric = self._resolve_search() - factory = ctx.require("factory") - - optimizer.optimize( + # Every read of the context happens here rather than in the + # shared helper: the contract audit parses this file, so a + # require moved out of it makes a declared key look unread. + model, best_parameters, plot_paths = self._search( model, x, y, optimizable_parameters, - goal_metric, + ctx.require("factory"), + # ctx.require already hands back an isolated copy of the + # stored parameter tree, so writing the best values into it + # cannot touch the Run row it came from. + ctx.require("model_parameters"), + self.config["run_id"], + self.config["artifact_prefix"], self._score_one_trial, ) - model = optimizer.get_model() - best_params = optimizer.get_best_params() - - self._assert_model_keeps_its_runtime_state(model) - - # ctx.require already hands back an isolated copy of the - # stored parameter tree, so update_parameters is free to - # mutate it without touching the Run row it came from. - old_parameters = ctx.require("model_parameters") - ctx.put_ref( - "best_parameters", - factory.update_parameters(old_parameters, best_params), - ) - - # Resolved here and not at the top of the method: the runs - # directory is only needed to name the plots a search produces, - # so a fit without one has no reason to require it of whatever - # is running it. - import os - import pickle - - from kink import di - - from DashAI.back.core.artifacts import normalize_artifacts - - config = di["config"] - - trials = optimizer.get_trials_values() - plot_filenames, plots = optimizer.create_plots( - trials, - run_id, - n_params=len(optimizable_parameters), - goal_metric=goal_metric, - artifact_prefix=self.config["artifact_prefix"], - ) - normalized_plots = normalize_artifacts(plots) - for filename, plot in zip( - plot_filenames, normalized_plots, strict=False - ): - plot_path = os.path.join(config["RUNS_PATH"], filename) - with open(plot_path, "wb") as file: - pickle.dump(plot, file) - plot_paths.append(plot_path) + ctx.put_ref("best_parameters", best_parameters) except Exception as e: log.exception(e) raise JobError( @@ -310,6 +114,10 @@ def execute(self, ctx: ExecutionContext) -> None: ctx.put("model", model) ctx.put_ref("plot_paths", plot_paths) + def _fit_kept_model(self, model, x, y) -> None: + """Fit the model that gets kept, on the partitions it was given.""" + self._fit(model, x, y, self.config.get("validation_during_fit", True)) + def _score_one_trial(self, model, x, y, metric) -> float: """Fit the model once and score it, for one point of the search. @@ -345,9 +153,7 @@ def _score_one_trial(self, model, x, y, metric) -> float: float The score on the validation partition, which is the objective. """ - from DashAI.back.core.enums.metrics import LevelEnum - - self._fit(model, x, y) + self._fit_kept_model(model, x, y) for name in self.config.get("trial_splits", TRIAL_SPLITS): model.calculate_metrics(split=SplitEnum[name], level=LevelEnum.TRIAL) @@ -355,52 +161,3 @@ def _score_one_trial(self, model, x, y, metric) -> float: predictions = model.predict(x["validation"]) expected = model.prepare_output(y["validation"], is_fit=False) return metric.score(expected, predictions) - - def _fit(self, model, x, y) -> None: - """Fit the model on the training partition of the data it was given. - - Whether the validation partition goes with it is a policy and not a - shape: a model uses it to watch the fit and stop early, which is what - an ordinary holdout run wants, and which is exactly wrong when that - same partition is what the fit will be scored on -- a fold is scored on - the rows it held back, so handing them over would be measuring the fit - on data it was allowed to watch. - - The key may also simply not be there. The trailing entry a fold - splitter produces holds the pooled rows and the reserved ones and has - no validation partition at all, so there is nothing to hand over even - where the policy would allow it. - """ - # ``.get`` with the schema's own placeholder, the way the other units - # read a declared optional field: a caller that builds this unit by - # hand -- a job, a test -- should not have to name a policy it is happy - # to leave alone, and the ordinary answer is the ordinary holdout one. - fit_with_validation = self.config.get("validation_during_fit", True) - if fit_with_validation and "validation" in x: - model.train(x["train"], y["train"], x["validation"], y["validation"]) - else: - model.train(x["train"], y["train"]) - - @staticmethod - def _assert_model_keeps_its_runtime_state(model) -> None: - """Fail loudly if the optimizer returned a model that cannot be scored. - - ``ModelFactory`` attaches the data splits and the metric classes to the - model instance, and optimizers are expected to return that same - instance. If one ever returns a fresh object instead, scoring it finds - nothing to score and the caller ends up with no metrics rather than an - error. - - The check is on the data, not on the run id. Keying it to ``run_id`` - made it a no-op for every caller that has no run -- a pipeline, where - ``run_id`` is always None -- which is exactly the caller with no other - signal that anything went wrong: it would finish with an empty metrics - artifact. What both callers need is the same, so this asks for that - instead. - """ - if getattr(model, "x_data", None) is None: - raise JobError( - "The optimizer returned a model detached from its data: metrics " - "could not be computed for it. Optimizers must return the same " - "model instance they received." - ) diff --git a/DashAI/back/units/fit_scope.py b/DashAI/back/units/fit_scope.py new file mode 100644 index 000000000..aa6fdba1d --- /dev/null +++ b/DashAI/back/units/fit_scope.py @@ -0,0 +1,354 @@ +"""Shared body of the units that fit a model, with or without a search. + +``FitModelUnit`` fits one set of partitions; ``FitModelOverFoldsUnit`` fits a +list of them. Everything around the fit is the same for both -- resolving the +optimizer and the goal metric, running the search, checking that the optimizer +gave back the model it was handed, recording the best parameters, and writing +the plots the search produces -- and it is exactly the sort of thing that ends +up written twice and then drifting. + +What differs is only the objective the search measures and what happens once it +is over, so those stay in the units. + +Named ``ModelFitScopeMixin`` rather than ``BaseSomething``: the registry derives +a component's type by walking its ``__mro__`` for a class whose name contains +"Base" and that declares ``TYPE``, and demands exactly one. A shared parent +called ``Base*`` would be a second candidate and would break the registration of +every unit that inherited it. +""" + +import logging +from typing import TYPE_CHECKING, List, Tuple + +from DashAI.back.core.schema_fields import ( + bool_field, + component_field, + enum_field, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError + +if TYPE_CHECKING: + from DashAI.back.optimizers.base_optimizer import BaseOptimizer + +log = logging.getLogger(__name__) + +#: A trial may score the partition it fitted on and the one it is measured +#: against, and nothing else. The test partition is deliberately absent: +#: scoring it once per trial would let the search see it, and a model chosen +#: with the test set in view has no honest score left to report. +TRIAL_SPLITS = ["TRAIN", "VALIDATION"] + + +def optimizer_field(): + return schema_field( + component_field(parent="BaseOptimizer"), + placeholder={"component": "OptunaOptimizer", "params": {}}, + description=MultilingualString( + en="Optimizer used to search for hyperparameters, along with its own " + "configuration. Only used when the model declares optimizable " + "parameters.", + es="Optimizador usado para buscar hiperparámetros, junto con su propia " + "configuración. Solo se usa cuando el modelo declara parámetros " + "optimizables.", + pt="Otimizador usado para procurar hiperparâmetros, junto com a sua " + "própria configuração. Só é usado quando o modelo declara parâmetros " + "otimizáveis.", + de="Optimierer für die Hyperparametersuche samt seiner eigenen " + "Konfiguration. Wird nur verwendet, wenn das Modell optimierbare " + "Parameter deklariert.", + zh="用于搜索超参数的优化器及其自身配置。仅当模型声明了可优化参数时使用。", + ), + alias=MultilingualString( + en="Optimizer", + es="Optimizador", + pt="Otimizador", + de="Optimierer", + zh="优化器", + ), + ) + + +def goal_metric_field(): + return schema_field( + string_field(), + placeholder="Accuracy", + description=MultilingualString( + en="Metric the hyperparameter search optimizes.", + es="Métrica que optimiza la búsqueda de hiperparámetros.", + pt="Métrica que a procura de hiperparâmetros otimiza.", + de="Metrik, die die Hyperparametersuche optimiert.", + zh="超参数搜索所优化的指标。", + ), + alias=MultilingualString( + en="Goal metric", + es="Métrica objetivo", + pt="Métrica objetivo", + de="Zielmetrik", + zh="目标指标", + ), + ) + + +def trial_splits_field(): + return schema_field( + list_field(enum_field(enum=TRIAL_SPLITS)), + placeholder=TRIAL_SPLITS, + description=MultilingualString( + en="Partitions each trial of the search records a score for. A " + "partition the model is not meant to be judged on belongs out of " + "this list even when metrics are configured for it.", + es="Particiones para las que cada intento de la búsqueda registra " + "un puntaje. Una partición sobre la que el modelo no debe juzgarse " + "no va en esta lista aunque tenga métricas configuradas.", + pt="Partições para as quais cada tentativa da procura regista uma " + "pontuação. Uma partição sobre a qual o modelo não deve ser julgado " + "fica fora desta lista mesmo que tenha métricas configuradas.", + de="Partitionen, für die jeder Versuch der Suche einen Wert " + "festhält. Eine Partition, nach der das Modell nicht beurteilt " + "werden soll, gehört nicht in diese Liste, auch wenn Metriken für " + "sie konfiguriert sind.", + zh="搜索的每次试验为其记录分数的分区。不应据以评判模型的分区不列入此处，" + "即使已为其配置了指标。", + ), + alias=MultilingualString( + en="Trial splits", + es="Particiones por intento", + pt="Partições por tentativa", + de="Versuchspartitionen", + zh="试验分区", + ), + ) + + +def validation_during_fit_field(): + return schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether the validation partition is handed to the model while " + "fitting. Models use it to watch training and stop early. Turn it " + "off when the same partition is what the fit will be scored on.", + es="Si la particion de validacion se entrega al modelo durante el " + "ajuste. Los modelos la usan para vigilar el entrenamiento y " + "detenerse antes. Desactivar cuando esa misma particion es sobre " + "la que se va a evaluar el ajuste.", + pt="Se a particao de validacao e entregue ao modelo durante o " + "ajuste. Os modelos usam-na para acompanhar o treino e parar mais " + "cedo. Desative quando essa mesma particao for aquela sobre a qual " + "o ajuste sera avaliado.", + de="Ob die Validierungspartition dem Modell beim Fitten uebergeben " + "wird. Modelle nutzen sie, um das Training zu beobachten und frueh " + "abzubrechen. Abschalten, wenn genau diese Partition den Fit " + "bewerten soll.", + zh="拟合时是否将验证分区交给模型。", + ), + alias=MultilingualString( + en="Validate while fitting", + es="Validar durante el ajuste", + pt="Validar durante o ajuste", + de="Beim Fitten validieren", + zh="拟合时验证", + ), + ) + + +class ModelFitScopeMixin: + """Everything around a fit that does not depend on how the data is shaped. + + Takes and returns plain values and never touches the execution context. + That is not tidiness: the contract audit parses each unit's own source, so + a ``ctx.require`` moved in here makes a declared key look unread and a + ``ctx.put`` makes a broken ``PROVIDES`` pass. Reading and publishing stay + in the units, and so does reading the runtime parameters, which the audit + checks are used by whoever declares them. + """ + + #: Declared here so a reader can see what state an instance carries; each + #: unit sets them in its own ``__init__``, because ``BaseUnit.__init__`` + #: comes first in the MRO and does not chain. + _optimizer = None + _goal_metric = None + + def _resolve_search(self): + """Resolve the optimizer and the goal metric, memoized on this unit. + + Kept on the instance rather than in the context on purpose. These are + this unit's own state, not something it hands to another unit: two + fitting units sharing a context -- a graph with two training nodes -- + would otherwise overwrite each other's optimizer, and the second one + would silently run the first one's. + """ + if self._optimizer is not None: + return self._optimizer, self._goal_metric + + from kink import di + + component_registry = di["component_registry"] + goal_metric_name: str = self.config["goal_metric"] + optimizer_name: str = self.config["optimizer"]["component"] + + try: + # The whole registry entry, not the class: the optimizer reads + # metadata["maximize"] from it to pick a direction. + goal_metric = component_registry[goal_metric_name] + except Exception as e: + log.exception(e) + raise JobError( + f"Metric is not compatible with the Task. {e}", + ) from e + + try: + optimizer_class = component_registry[optimizer_name]["class"] + optimizer: "BaseOptimizer" = optimizer_class( + **self.config["optimizer"]["params"] + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instantiating optimizer {optimizer_name}, {e}", + ) from e + + self._goal_metric = goal_metric + self._optimizer = optimizer + return optimizer, goal_metric + + def _validate_search(self, optimizable_parameters) -> None: + """Refuse an impossible search before anything observable happens. + + Handed the value rather than the context: the caller reads it with + ``ctx.require`` and not ``ctx.get``, because an absent key means the + model has not been built yet -- a call-order mistake, not "there is + nothing to optimize". Only an empty value, the key present and the + model declaring none, skips the checks below, so no registry lookup is + needed either. + """ + if not optimizable_parameters: + return + + self._resolve_search() + + def _fit(self, model, x, y, with_validation: bool) -> None: + """Fit the model on the training partition of the data it was given. + + Whether the validation partition goes with it is a policy and not a + shape: a model uses it to watch the fit and stop early, which is what + an ordinary holdout run wants, and which is exactly wrong when that + same partition is what the fit will be scored on -- a fold is scored on + the rows it held back, so handing them over would be measuring the fit + on data it was allowed to watch. + + The key may also simply not be there. The trailing entry a fold + splitter produces holds the pooled rows and the reserved ones and has + no validation partition at all, so there is nothing to hand over even + where the policy would allow it. + """ + if with_validation and "validation" in x: + model.train(x["train"], y["train"], x["validation"], y["validation"]) + else: + model.train(x["train"], y["train"]) + + def _search( + self, + model, + x, + y, + optimizable_parameters, + factory, + old_parameters, + run_id, + artifact_prefix, + objective, + ) -> Tuple[object, dict, List[str]]: + """Run the hyperparameter search and report what it produced. + + Parameters + ---------- + model : BaseModel + The instance the search sets its parameters on. + x, y : object + Whatever the objective knows how to fit and score -- one set of + partitions, or a list of them. The optimizer passes it through + without looking at it, which is what lets one search serve both. + optimizable_parameters : list + The search space, as ModelFactory built it. + factory : ModelFactory + Used to write the values found back into the parameter tree. + old_parameters : dict + The tree they are written into. The caller is expected to hand over + a copy it owns; this mutates nothing else. + run_id, artifact_prefix : object + What the plots are named after. + objective : callable + ``(model, x, y, metric) -> float``: what the search measures. + + Returns + ------- + tuple + The fitted model, the parameter tree with the best values in it, + and the paths of the plots the search produced. + """ + import os + import pickle + + from kink import di + + from DashAI.back.core.artifacts import normalize_artifacts + + # Memoized: validate() resolved these already, and resolving again + # here would be the same lookup. + optimizer, goal_metric = self._resolve_search() + + optimizer.optimize(model, x, y, optimizable_parameters, goal_metric, objective) + model = optimizer.get_model() + best_params = optimizer.get_best_params() + + self._assert_model_keeps_its_runtime_state(model) + + best_parameters = factory.update_parameters(old_parameters, best_params) + + config = di["config"] + plot_paths: List[str] = [] + trials = optimizer.get_trials_values() + plot_filenames, plots = optimizer.create_plots( + trials, + run_id, + n_params=len(optimizable_parameters), + goal_metric=goal_metric, + artifact_prefix=artifact_prefix, + ) + normalized_plots = normalize_artifacts(plots) + for filename, plot in zip(plot_filenames, normalized_plots, strict=False): + plot_path = os.path.join(config["RUNS_PATH"], filename) + with open(plot_path, "wb") as file: + pickle.dump(plot, file) + plot_paths.append(plot_path) + + return model, best_parameters, plot_paths + + @staticmethod + def _assert_model_keeps_its_runtime_state(model) -> None: + """Fail loudly if the optimizer returned a model that cannot be scored. + + ``ModelFactory`` attaches the metric classes to the model instance and + whoever fits it points it at its data, and optimizers are expected to + return that same instance. If one ever returns a fresh object instead, + scoring it finds nothing to score and the caller ends up with no + metrics rather than an error. + + The check is on the data, not on the run id. Keying it to ``run_id`` + made it a no-op for every caller that has no run -- a pipeline, where + ``run_id`` is always None -- which is exactly the caller with no other + signal that anything went wrong: it would finish with an empty metrics + artifact. What both callers need is the same, so this asks for that + instead. + """ + if getattr(model, "x_data", None) is None: + raise JobError( + "The optimizer returned a model detached from its data: metrics " + "could not be computed for it. Optimizers must return the same " + "model instance they received." + ) From f8a3fc7bf2eaf70915e1cdde674111d5a8bef3b5 Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 19:19:36 -0300 Subject: [PATCH 24/28] Train a cross-validated run through the units, unless it is nested FitModelOverFoldsUnit is FitModelUnit's sibling: it takes a list of partition sets instead of one, which is a different REQUIRES and so a different unit. Everything around the fit is the shared mixin; what is written here is the objective the search measures -- the whole fold loop, so one trial costs k fits -- and what happens once the search is over. ModelJob composes it for a fold run that is not nested. Nested cross-validation still trains through the strategy: its inner splitter is a required component field, and a component field cannot be made optional without leaving the user without a selector, so it is a further sibling rather than a flag on this one. Three decisions worth naming. The per-fold scores are published rather than aggregated in the unit. A summary row carries a standard deviation, and a unit may not write domain rows -- the one sanctioned write in the domain layer has nowhere to put one. So the unit hands the numbers over and the job does the arithmetic and the writing, where every other row it persists is written. A single fold gets a deviation of zero rather than none, because none is what the reserved-rows measurement carries and the two say different things. A trial records one row per split holding the mean over its folds, not one per fold: the folds of a trial measure a hyperparameter setting rather than the model that gets kept, and recording each would bury the rows that describe it. That write is guarded on the run, which _save_metrics does not guard for itself -- a caller with no run would write rows against a foreign key pointing at nothing, and they insert without complaint because nothing enforces it. Scoring the reserved rows is not the unit's. It is an ordinary LAST metric, so it is EvaluateModelUnit, the same one a holdout run uses, and whether there is anything to score is the caller's to know: a session that reserved nothing leaves that partition empty rather than absent. calculate_metrics now returns what it wrote, so a caller that wants both the row and the number scores the split once instead of twice. The assertion that the optimizer gave back the model it was handed caught a real gap while this was written: with the data attached at fit time rather than at build time, nothing had pointed the model at anything during a fold search. It does now, per fold, the same as the scoring loop. 1001 passed across units, dag, spike, api and evaluation; both nets unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/initial_components.py | 2 + DashAI/back/job/model_job.py | 111 ++++++- DashAI/back/models/base_model.py | 13 +- .../back/units/fit_model_over_folds_unit.py | 276 ++++++++++++++++++ tests/back/api/test_units_api.py | 3 + 5 files changed, 399 insertions(+), 6 deletions(-) create mode 100644 DashAI/back/units/fit_model_over_folds_unit.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 2339426ab..daad62e08 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -524,6 +524,7 @@ ) from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_converter_unit import FitConverterUnit +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.generate_global_explanation_unit import ( GenerateGlobalExplanationUnit, @@ -732,6 +733,7 @@ def get_initial_components(): PrepareAndFoldUnit, BuildModelUnit, FitModelUnit, + FitModelOverFoldsUnit, EvaluateModelUnit, EvaluateModelToArtifactUnit, SaveModelUnit, diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index 6475f4b08..da670daec 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -5,7 +5,13 @@ from sqlalchemy import exc from sqlalchemy.orm.attributes import flag_modified -from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.dependencies.database.models import ( + Dataset, + Metric, + ModelSession, + Run, +) from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy from DashAI.back.job.base_job import BaseJob, JobError from DashAI.back.optimizers.base_optimizer import BaseOptimizer @@ -13,6 +19,7 @@ from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.context import ExecutionContext from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.load_dataset_unit import LoadDatasetUnit from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit @@ -224,10 +231,47 @@ def run( self.report_progress(0.85, "Computing metrics") EvaluateModelUnit(run_id=run_id, splits=scored_splits)(ctx) + elif not run.nested: + fit_folds = FitModelOverFoldsUnit( + optimizer={ + "component": run.optimizer_name, + "params": run.optimizer_parameters, + }, + goal_metric=run.goal_metric, + run_id=run_id, + artifact_prefix=str(run_id), + scored_splits=[ + name for name in scored_splits if name != "TEST" + ], + ) + fit_folds(ctx) + + plot_paths = ctx.require("plot_paths") + if ctx.has("best_parameters"): + run.parameters = ctx.get("best_parameters") + flag_modified(run, "parameters") + db.commit() + + self.report_progress(0.85, "Computing metrics") + self._aggregate_fold_metrics( + db, run_id, ctx.require("fold_metrics") + ) + + # The rows the session reserved are the only ones no + # fold and no trial ever saw, so they are the only + # honest estimate left once a model is picked out of a + # comparison table -- and scoring them is an ordinary + # LAST metric, so it is the same unit a holdout run + # uses. Whether there is anything to score is the + # caller's to know: a session that reserved nothing + # leaves that partition empty rather than absent. + if len(x[-1]["test"]) > 0: + EvaluateModelUnit(run_id=run_id, splits=["TEST"])(ctx) else: - # Fold runs still train through the strategy. Their loop - # is the next piece to move; everything before and after - # it is already the units'. + # Nested cross-validation still trains through the + # strategy: its inner splitter is a required component + # field, so it is a further sibling unit rather than a + # flag on the one above, and it is not written yet. evaluation_estrategy: BaseEvaluationStrategy = strategy_class( factory=ctx.require("factory"), optimizer=preparation_results["optimizer"], @@ -295,6 +339,65 @@ def run( ctx.clear_cache() gc.collect() + @staticmethod + def _aggregate_fold_metrics(db, run_id: int, fold_metrics: Dict[str, Any]) -> None: + """Summarise the per-fold scores into one row per split and metric. + + The unit that fitted the folds publishes their scores rather than + aggregating them, because a summary row carries a standard deviation + and a unit may not write domain rows -- the one sanctioned write in the + domain layer has nowhere to put one. So the arithmetic and the writing + happen here, where every other row this job persists is written. + + A single fold gets a deviation of zero rather than none: none is what + the reserved-rows measurement carries, and the two say different + things -- "one fold, so nothing varied" against "not a summary at all". + + Parameters + ---------- + db : Session + The session this job is already holding. + run_id : int + The run the rows belong to. + fold_metrics : dict + ``{split name: {metric name: [one score per fold]}}``. + """ + import numpy as np + + for split_name, by_metric in fold_metrics.items(): + for metric_name, values in by_metric.items(): + if not values: + continue + existing = ( + db.query(Metric) + .filter_by( + run_id=run_id, + split=SplitEnum[split_name], + level=LevelEnum.LAST, + name=metric_name, + ) + .first() + ) + mean = float(np.mean(values)) + deviation = float(np.std(values)) if len(values) > 1 else 0.0 + + if existing: + existing.value = mean + existing.std_value = deviation + else: + db.add( + Metric( + run_id=run_id, + split=SplitEnum[split_name], + level=LevelEnum.LAST, + name=metric_name, + value=mean, + std_value=deviation, + step=0, + ) + ) + db.commit() + def _prepare_dataset_and_components( self, run_id: int, db, component_registry ) -> Dict[str, Any]: diff --git a/DashAI/back/models/base_model.py b/DashAI/back/models/base_model.py index c5d1e7f6b..144f44ce5 100644 --- a/DashAI/back/models/base_model.py +++ b/DashAI/back/models/base_model.py @@ -349,6 +349,13 @@ def calculate_metrics( labels stored in the model for the given split are used. Defaults to None. + Returns + ------- + Dict[str, float] or None + What was written, so a caller that also wants the numbers does not + have to score the same split twice. ``None`` when nothing was + written: no run to write against, or nothing to score. + Notes ----- A metric row is keyed by the run it belongs to, so a model with no run @@ -366,11 +373,11 @@ def calculate_metrics( # checked for metrics first. Treating "no attribute" as "no run" keeps # that path working and is the same answer for any caller that has one. if not getattr(self, "run_id", None): - return + return None results = self.compute_metrics(split=split, x_data=x_data, y_data=y_data) if results is None: - return + return None # Save to database self._save_metrics( @@ -392,6 +399,8 @@ def calculate_metrics( ): self._epoch_reporter(results, log_index) + return results + def prepare_dataset( self, dataset: "DashAIDataset", is_fit: bool = False ) -> "DashAIDataset": diff --git a/DashAI/back/units/fit_model_over_folds_unit.py b/DashAI/back/units/fit_model_over_folds_unit.py new file mode 100644 index 000000000..ab16420d2 --- /dev/null +++ b/DashAI/back/units/fit_model_over_folds_unit.py @@ -0,0 +1,276 @@ +"""Unit that fits a model across cross-validation folds.""" + +import logging + +import numpy as np + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema +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.fit_scope import ( + TRIAL_SPLITS, + ModelFitScopeMixin, + goal_metric_field, + optimizer_field, + trial_splits_field, +) + +log = logging.getLogger(__name__) + + +class FitModelOverFoldsSchema(BaseSchema): + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + scored_splits: trial_splits_field() # type: ignore + + +class FitModelOverFoldsUnit(BaseUnit, ModelFitScopeMixin): + """Fit a model once per fold, then once more on everything the folds used. + + The sibling of ``FitModelUnit``. It takes a list of partition sets rather + than one, which is a different ``REQUIRES`` and therefore a different unit; + everything around the fit is shared between them. + + Three things happen here, in this order, and the order is the point: + + 1. If the model declares optimizable parameters, the search runs first. + Its objective is the whole fold loop, so one trial costs k fits. That is + what cross-validation buys and what it costs, and it is why the + objective had to become something the caller supplies rather than + something the optimizer does. + 2. Every fold is fitted and scored, and a ``FOLD`` metric row is written + for it. Those rows are what the fold charts and the statistical tests + read, and they are indexed from zero without gaps -- repeated + cross-validation buckets them by integer division, so a gap silently + moves a fold into the wrong repetition. + 3. The model that gets kept is fitted on every row the folds could use, + and left pointing at that partition set. Scoring it against the rows the + session reserved is not done here: that is a ``LAST`` metric like any + other, so it is ``EvaluateModelUnit``'s, the same unit a holdout run + uses -- and the caller is the one that knows whether anything was + reserved to score against. + + **The per-fold scores are published rather than aggregated here.** Turning + them into a mean and a deviation means writing ``Metric`` rows that carry a + ``std_value``, and a unit may not write domain rows -- the one sanctioned + write in the domain layer, ``BaseModel._save_metrics``, has nowhere to put + a deviation. Whoever asked for the fit aggregates and persists. + + Nested cross-validation is not here either: its inner splitter is a + required component field, and a component field cannot be made optional + without leaving the user without a selector, so it is a further sibling + rather than a flag on this one. + """ + + SCHEMA = FitModelOverFoldsSchema + + REQUIRES = ( + "model", + "factory", + "optimizable_parameters", + "model_parameters", + "x_folds", + "y_folds", + ) + PROVIDES = ("model", "plot_paths", "fold_metrics") + RUNTIME_PARAMS = ("run_id", "artifact_prefix") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._optimizer = None + self._goal_metric = None + + def validate(self, ctx: ExecutionContext) -> None: + # ctx.require, not ctx.get: an absent key means BuildModelUnit has not + # run yet, which is a call-order mistake rather than a model with + # nothing to optimize. + self._validate_search(ctx.require("optimizable_parameters")) + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + x_folds = ctx.require("x_folds") + y_folds = ctx.require("y_folds") + optimizable_parameters = ctx.require("optimizable_parameters") + + plot_paths = [] + try: + if optimizable_parameters: + # Every read of the context happens in this file rather than in + # the shared helper: the contract audit parses it, so a require + # moved out makes a declared key look unread. + model, best_parameters, plot_paths = self._search( + model, + x_folds, + y_folds, + optimizable_parameters, + ctx.require("factory"), + ctx.require("model_parameters"), + self.config["run_id"], + self.config["artifact_prefix"], + self._score_one_trial, + ) + ctx.put_ref("best_parameters", best_parameters) + + fold_metrics = self._score_every_fold(model, x_folds, y_folds) + + # The model that gets kept is fitted last, on every row the folds + # could use, and is left pointing at that partition set -- which is + # what whoever scores it afterwards will be scoring. + model.x_data = x_folds[-1] + model.y_data = y_folds[-1] + self._fit_kept_model(model, x_folds[-1], y_folds[-1]) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + + ctx.put("model", model) + ctx.put_ref("plot_paths", plot_paths) + ctx.put_ref("fold_metrics", fold_metrics) + + # ----------------------------------------------------------------- # + # The three things, one method each + # ----------------------------------------------------------------- # + + def _score_one_trial(self, model, x_folds, y_folds, metric) -> float: + """Fit and score every fold, and return the mean, for one trial. + + This is what the optimizer measures, and it is the whole difference + between this unit and its sibling: there, a trial is one fit; here it + is k of them. The optimizer passes the partitions through without + looking at them, so it does not have to know which. + + One row per split is written for the trial, holding the mean over its + folds -- not one row per fold. The folds of a trial measure a + hyperparameter setting rather than the model that gets kept, so + recording each of them would bury the rows that describe it, and the + live chart that watches a search wants one point per trial anyway. + + The objective is scored directly rather than read out of those means, + because the goal metric is not necessarily one of the metrics the + session configured for the validation partition. + + Returns + ------- + float + The mean of the goal metric over the folds, which is the objective. + """ + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + scores = [] + accumulated: dict = {} + + for x_fold, y_fold in self._folds(x_folds, y_folds): + # Pointed at the fold it is about to be fitted on, the same as in + # the scoring loop. Nothing here reads it back, but the model is + # expected to carry the data it was last fitted on -- that is what + # the check after the search asserts, and what anything scoring it + # afterwards relies on. + model.x_data = x_fold + model.y_data = y_fold + self._fit_kept_model(model, x_fold, y_fold) + predictions = model.predict(x_fold["validation"]) + expected = model.prepare_output(y_fold["validation"], is_fit=False) + scores.append(metric.score(expected, predictions)) + + for name in scored_splits: + results = model.compute_metrics(split=SplitEnum[name]) + for metric_name, value in (results or {}).items(): + accumulated.setdefault(name, {}).setdefault(metric_name, []).append( + value + ) + + self._record_trial(model, accumulated) + return float(np.mean(scores)) + + @staticmethod + def _record_trial(model, accumulated: dict) -> None: + """Write one metric row per split for this trial, holding its mean. + + Guarded on the run, which ``_save_metrics`` does not do for itself: a + caller with no run -- a pipeline -- would otherwise write rows against + a foreign key pointing at nothing, and they insert without complaint + because the database has no enforcement to refuse them. That is the + failure mode a metrics unit already exists to avoid, so it is refused + here too rather than repeated. + """ + if not getattr(model, "run_id", None): + return + + for split_name, by_metric in accumulated.items(): + averaged = { + metric_name: float(np.mean(values)) + for metric_name, values in by_metric.items() + } + if averaged: + model._save_metrics( + results=averaged, + split=SplitEnum[split_name], + level=LevelEnum.TRIAL, + ) + + def _score_every_fold(self, model, x_folds, y_folds) -> dict: + """Fit and score each fold, writing its rows and keeping its numbers. + + Returns + ------- + dict + ``{split name: [one score per fold]}``, in fold order, for whoever + aggregates them. A split with no metrics configured is absent + rather than present and empty: the two are different statements. + """ + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + fold_metrics: dict = {} + + for index, (x_fold, y_fold) in enumerate(self._folds(x_folds, y_folds)): + # The metric methods read the data off the instance to decide what + # they are scoring, so it follows the iteration. + model.x_data = x_fold + model.y_data = y_fold + self._fit_kept_model(model, x_fold, y_fold) + + for name in scored_splits: + split = SplitEnum[name] + # calculate_metrics hands back what it wrote, so the split is + # scored once rather than once for the row and once for the + # number. It writes nothing, and returns nothing, when there is + # no run to write against -- a pipeline -- and the numbers are + # wanted there too, hence the fallback. + results = model.calculate_metrics( + split=split, level=LevelEnum.FOLD, fold_index=index + ) + if results is None: + results = model.compute_metrics(split=split) + if not results: + continue + for metric_name, value in results.items(): + fold_metrics.setdefault(name, {}).setdefault( + metric_name, [] + ).append(value) + + return fold_metrics + + def _fit_kept_model(self, model, x, y) -> None: + """Fit on the training partition, never handing over the validation one. + + A fold is scored on the rows it held back from its own training, so + letting the fit watch them would measure it on data it was allowed to + see -- and nothing raises when that happens, the score simply comes out + better than the model deserves. Unlike the holdout sibling this is not + configurable, because there is no reading of a fold under which it + would be right. + """ + self._fit(model, x, y, with_validation=False) + + @staticmethod + def _folds(x_folds, y_folds): + """Pair up the folds, leaving out the entry that is not one. + + A fold splitter returns one entry per fold plus a trailing entry + holding every row the folds could use and the rows reserved from them. + That last one fits the model that gets kept; it is not a fold and is + never scored as one. + """ + return zip(x_folds[:-1], y_folds[:-1], strict=True) diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index fea24053a..7c33c0497 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -9,6 +9,7 @@ "PrepareAndFoldUnit", "BuildModelUnit", "FitModelUnit", + "FitModelOverFoldsUnit", "EvaluateModelUnit", "EvaluateModelToArtifactUnit", "SaveModelUnit", @@ -231,6 +232,8 @@ def get_unit_classes(client: TestClient): ("BuildModelUnit", "run_id"), ("EvaluateModelUnit", "run_id"), ("FitModelUnit", "run_id"), + ("FitModelOverFoldsUnit", "run_id"), + ("FitModelOverFoldsUnit", "artifact_prefix"), ("FitModelUnit", "artifact_prefix"), ("SaveModelUnit", "artifact_prefix"), } From 77dbca1a3c065d921ce28b82d378e83be48f7f4f Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 19:41:25 -0300 Subject: [PATCH 25/28] Measure a nested run through the units too, and stop calling the strategy FitModelOverNestedFoldsUnit is FitModelOverFoldsUnit plus one measurement taken before it. Inheritance rather than a shared mixin because that is the actual relationship: everything the sibling does still happens, and this adds a step in front. Two units and not one with a flag because its inner splitter is a required component field, and a component field cannot be made optional -- the front reads `parent` straight off the property, and an anyOf buries it where it does not look. What the nested loop is for, since the code alone does not say it: in an ordinary cross-validated search the same folds choose the hyperparameters and report the score, so the score is optimistic by however much the search managed to fit them. The nested loop measures that honestly -- for each outer fold a search runs on folds carved out of that fold's training rows alone, and what it chooses is scored on the outer fold's validation rows, which it never saw. What it does not do is choose the hyperparameters: each outer fold picks its own and they generally differ, so there is no single model to keep out of that loop. The ordinary search still runs afterwards and produces the model that gets saved. The nested numbers describe the procedure, not the artifact, which is why they are kept at their own level -- LAST_OUTER against LAST -- and why the inner trials record nothing at all. The two fold branches in the job became one, choosing a unit rather than repeating a body. **The evaluation strategies are no longer called.** The job reads SCORED_SPLITS and KIND off the class it resolves and never touches execute() on any path. The classes are now what they always were underneath -- a declaration of how a run is carved and what it records -- and emptying them of the code that is now unreachable is the last piece. 1014 passed across units, dag, spike, api and evaluation, both nets unchanged, plus thirteen contract tests for the fold unit built on a hand-made context: the end-to-end net runs it inside a real job, which cannot show what it reads, promises and refuses on its own. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/initial_components.py | 4 + DashAI/back/job/model_job.py | 85 +++-- .../back/units/fit_model_over_folds_unit.py | 6 +- .../units/fit_model_over_nested_folds_unit.py | 195 +++++++++++ tests/back/api/test_units_api.py | 3 + .../units/test_fit_model_over_folds_unit.py | 310 ++++++++++++++++++ 6 files changed, 567 insertions(+), 36 deletions(-) create mode 100644 DashAI/back/units/fit_model_over_nested_folds_unit.py create mode 100644 tests/back/units/test_fit_model_over_folds_unit.py diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index daad62e08..fa358c163 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -525,6 +525,9 @@ from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_converter_unit import FitConverterUnit from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_model_over_nested_folds_unit import ( + FitModelOverNestedFoldsUnit, +) from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.generate_global_explanation_unit import ( GenerateGlobalExplanationUnit, @@ -734,6 +737,7 @@ def get_initial_components(): BuildModelUnit, FitModelUnit, FitModelOverFoldsUnit, + FitModelOverNestedFoldsUnit, EvaluateModelUnit, EvaluateModelToArtifactUnit, SaveModelUnit, diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index da670daec..c105dee89 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -12,7 +12,6 @@ ModelSession, Run, ) -from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy from DashAI.back.job.base_job import BaseJob, JobError from DashAI.back.optimizers.base_optimizer import BaseOptimizer from DashAI.back.splitters.splits_payload import normalize_splits_payload @@ -20,6 +19,9 @@ from DashAI.back.units.context import ExecutionContext from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_model_over_nested_folds_unit import ( + FitModelOverNestedFoldsUnit, +) from DashAI.back.units.fit_model_unit import FitModelUnit from DashAI.back.units.load_dataset_unit import LoadDatasetUnit from DashAI.back.units.prepare_and_fold_unit import PrepareAndFoldUnit @@ -164,8 +166,10 @@ def run( # the shape they publish for it. prepare(ctx) + # Only the partitions are needed here, and only to ask + # whether the session reserved any rows: the units read + # what they work on from the context themselves. x = ctx.get("x") if ctx.has("x") else ctx.require("x_folds") - y = ctx.get("y") if ctx.has("y") else ctx.require("y_folds") # save the obtained splits into the database run.split_indexes = json.dumps(ctx.require("split_indexes")) @@ -231,19 +235,32 @@ def run( self.report_progress(0.85, "Computing metrics") EvaluateModelUnit(run_id=run_id, splits=scored_splits)(ctx) - elif not run.nested: - fit_folds = FitModelOverFoldsUnit( - optimizer={ + else: + # Two units and not one with a flag: the nested one + # takes a required component field for its inner + # splitter, and a component field cannot be made + # optional without leaving the user without a selector. + fold_config = { + "optimizer": { "component": run.optimizer_name, "params": run.optimizer_parameters, }, - goal_metric=run.goal_metric, - run_id=run_id, - artifact_prefix=str(run_id), - scored_splits=[ + "goal_metric": run.goal_metric, + "run_id": run_id, + "artifact_prefix": str(run_id), + "scored_splits": [ name for name in scored_splits if name != "TEST" ], - ) + } + if run.nested: + fold_config["inner_splitter"] = { + "component": run.nested.get("splitter_name"), + "params": run.nested, + } + fit_folds = FitModelOverNestedFoldsUnit(**fold_config) + else: + fit_folds = FitModelOverFoldsUnit(**fold_config) + fit_folds(ctx) plot_paths = ctx.require("plot_paths") @@ -253,8 +270,20 @@ def run( db.commit() self.report_progress(0.85, "Computing metrics") + if ctx.has("outer_fold_metrics"): + # Kept at its own level: it answers a different + # question from the ordinary summary -- how the + # procedure does, rather than how this model does -- + # and the two would be indistinguishable side by + # side. + self._aggregate_fold_metrics( + db, + run_id, + ctx.get("outer_fold_metrics"), + LevelEnum.LAST_OUTER, + ) self._aggregate_fold_metrics( - db, run_id, ctx.require("fold_metrics") + db, run_id, ctx.require("fold_metrics"), LevelEnum.LAST ) # The rows the session reserved are the only ones no @@ -267,27 +296,6 @@ def run( # leaves that partition empty rather than absent. if len(x[-1]["test"]) > 0: EvaluateModelUnit(run_id=run_id, splits=["TEST"])(ctx) - else: - # Nested cross-validation still trains through the - # strategy: its inner splitter is a required component - # field, so it is a further sibling unit rather than a - # flag on the one above, and it is not written yet. - evaluation_estrategy: BaseEvaluationStrategy = strategy_class( - factory=ctx.require("factory"), - optimizer=preparation_results["optimizer"], - goal_metric=preparation_results["goal_metric"], - ) - evaluation_estrategy.set_progress_reporter(self.report_progress) - model, plot_paths = evaluation_estrategy.execute( - x=x, - y=y, - run=run, - db=db, - ) - # The strategy hands the model back rather than leaving - # it in the context, so the saving unit below serves - # both paths. - ctx.put("model", model) except Exception as e: log.exception(e) raise JobError( @@ -340,7 +348,9 @@ def run( gc.collect() @staticmethod - def _aggregate_fold_metrics(db, run_id: int, fold_metrics: Dict[str, Any]) -> None: + def _aggregate_fold_metrics( + db, run_id: int, fold_metrics: Dict[str, Any], level: LevelEnum + ) -> None: """Summarise the per-fold scores into one row per split and metric. The unit that fitted the folds publishes their scores rather than @@ -361,6 +371,11 @@ def _aggregate_fold_metrics(db, run_id: int, fold_metrics: Dict[str, Any]) -> No The run the rows belong to. fold_metrics : dict ``{split name: {metric name: [one score per fold]}}``. + level : LevelEnum + Where the summary goes. The ordinary fold scores summarise to + ``LAST``; the outer folds of a nested run summarise to + ``LAST_OUTER``, because they answer a different question and would + be indistinguishable from the first if they shared a level. """ import numpy as np @@ -373,7 +388,7 @@ def _aggregate_fold_metrics(db, run_id: int, fold_metrics: Dict[str, Any]) -> No .filter_by( run_id=run_id, split=SplitEnum[split_name], - level=LevelEnum.LAST, + level=level, name=metric_name, ) .first() @@ -389,7 +404,7 @@ def _aggregate_fold_metrics(db, run_id: int, fold_metrics: Dict[str, Any]) -> No Metric( run_id=run_id, split=SplitEnum[split_name], - level=LevelEnum.LAST, + level=level, name=metric_name, value=mean, std_value=deviation, diff --git a/DashAI/back/units/fit_model_over_folds_unit.py b/DashAI/back/units/fit_model_over_folds_unit.py index ab16420d2..662481002 100644 --- a/DashAI/back/units/fit_model_over_folds_unit.py +++ b/DashAI/back/units/fit_model_over_folds_unit.py @@ -136,6 +136,9 @@ def execute(self, ctx: ExecutionContext) -> None: # ----------------------------------------------------------------- # def _score_one_trial(self, model, x_folds, y_folds, metric) -> float: + return self._score_folds(model, x_folds, y_folds, metric, record=True) + + def _score_folds(self, model, x_folds, y_folds, metric, record: bool) -> float: """Fit and score every fold, and return the mean, for one trial. This is what the optimizer measures, and it is the whole difference @@ -182,7 +185,8 @@ def _score_one_trial(self, model, x_folds, y_folds, metric) -> float: value ) - self._record_trial(model, accumulated) + if record: + self._record_trial(model, accumulated) return float(np.mean(scores)) @staticmethod diff --git a/DashAI/back/units/fit_model_over_nested_folds_unit.py b/DashAI/back/units/fit_model_over_nested_folds_unit.py new file mode 100644 index 000000000..f55505da5 --- /dev/null +++ b/DashAI/back/units/fit_model_over_nested_folds_unit.py @@ -0,0 +1,195 @@ +"""Unit that measures a fold run with a search of its own inside each fold.""" + +import logging + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.core.schema_fields import BaseSchema +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit +from DashAI.back.units.fit_scope import ( + TRIAL_SPLITS, + goal_metric_field, + optimizer_field, + trial_splits_field, +) +from DashAI.back.units.splitter_scope import _splitter_field + +log = logging.getLogger(__name__) + + +def inner_splitter_field(): + """The splitter that carves an outer fold into folds of its own. + + Required rather than optional, which is what makes this a separate unit. + A component field wrapped in ``none_type`` is emitted as ``anyOf`` and the + front reads ``parent`` directly off the property, so an optional one leaves + the user with no selector at all -- the same wall that made the two + explainer units siblings. + """ + return _splitter_field( + parent="FoldSplitter", + placeholder={ + "component": "KFoldSplitter", + "params": {"n_splits": 3, "shuffle": True, "random_state": 42}, + }, + ) + + +class FitModelOverNestedFoldsSchema(BaseSchema): + optimizer: optimizer_field() # type: ignore + goal_metric: goal_metric_field() # type: ignore + scored_splits: trial_splits_field() # type: ignore + inner_splitter: inner_splitter_field() # type: ignore + + +class FitModelOverNestedFoldsUnit(FitModelOverFoldsUnit): + """Score every outer fold with a search that never saw it. + + ``FitModelOverFoldsUnit`` plus one measurement taken before it. The + relationship is inheritance rather than a shared mixin because that is what + it is: everything the sibling does still happens, and this adds a step in + front of it. + + **What the extra step is for.** In an ordinary cross-validated search, the + same folds choose the hyperparameters and report the score, so the reported + score is optimistic by however much the search managed to fit them. The + nested loop measures that honestly: for each outer fold, a search is run + from scratch on folds carved out of *that fold's training rows only*, and + the chosen model is then scored on the outer fold's validation rows, which + that search never saw. Those are the ``OUTER_FOLD`` rows. + + What it does **not** do is choose the hyperparameters. Each outer fold + picks its own, and they generally differ; there is no single model to keep + out of that loop. So the ordinary search still runs afterwards, over all + the folds, and it is what produces the model that gets saved. The nested + numbers are a statement about the procedure, not about the artifact. + + That also means the inner trials record nothing: their partitions belong to + one outer fold, and rows from them would sit alongside rows describing the + kept model as if they were comparable. + """ + + SCHEMA = FitModelOverNestedFoldsSchema + + PROVIDES = ("model", "plot_paths", "fold_metrics", "outer_fold_metrics") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._inner_splitter = None + + def _resolve_inner_splitter(self): + """Build the splitter that carves an outer fold, memoized on this unit.""" + if self._inner_splitter is not None: + return self._inner_splitter + + from kink import di + + chosen = self.config["inner_splitter"] + try: + splitter_class = di["component_registry"][chosen["component"]]["class"] + self._inner_splitter = splitter_class(splits_data=dict(chosen["params"])) + except Exception as e: + log.exception(e) + raise JobError( + f"Error configuring inner splitter for nested CV: {e}", + ) from e + return self._inner_splitter + + def execute(self, ctx: ExecutionContext) -> None: + model = ctx.require("model") + x_folds = ctx.require("x_folds") + y_folds = ctx.require("y_folds") + optimizable_parameters = ctx.require("optimizable_parameters") + + # Resolved outside the wrapper below so a splitter that cannot be built + # is reported as that, rather than as a training failure. + inner_splitter = self._resolve_inner_splitter() + + if optimizable_parameters: + try: + outer_fold_metrics = self._measure_every_outer_fold( + model, + x_folds, + y_folds, + inner_splitter, + optimizable_parameters, + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Model training failed {e}", + ) from e + else: + # Nothing to search means nothing for the nested loop to measure: + # every outer fold would choose the same parameters, which is what + # the ordinary loop already reports. + outer_fold_metrics = {} + + ctx.put_ref("outer_fold_metrics", outer_fold_metrics) + + # And then the ordinary run, which is what produces the kept model. + super().execute(ctx) + + def _measure_every_outer_fold( + self, model, x_folds, y_folds, inner_splitter, optimizable_parameters + ) -> dict: + """Search inside each outer fold, then score that fold with what it chose. + + Returns + ------- + dict + ``{split name: {metric name: [one score per outer fold]}}``, for + whoever aggregates them -- the same shape and the same reason as + the ordinary fold scores. + """ + optimizer, goal_metric = self._resolve_search() + scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) + outer_fold_metrics: dict = {} + + for index, (x_outer, y_outer) in enumerate(self._folds(x_folds, y_folds)): + # Carved out of this fold's training rows alone. Its validation + # rows are what the search will be judged on, so nothing drawn from + # them may reach it. + inner_x, inner_y, _ = inner_splitter.split( + x_outer["train"], y_outer["train"] + ) + + optimizer.optimize( + model, + inner_x, + inner_y, + optimizable_parameters, + goal_metric, + self._score_one_inner_trial, + ) + outer_model = optimizer.get_model() + + outer_model.x_data = x_outer + outer_model.y_data = y_outer + self._fit_kept_model(outer_model, x_outer, y_outer) + + for name in scored_splits: + split = SplitEnum[name] + results = outer_model.calculate_metrics( + split=split, level=LevelEnum.OUTER_FOLD, fold_index=index + ) + if results is None: + results = outer_model.compute_metrics(split=split) + if not results: + continue + for metric_name, value in results.items(): + outer_fold_metrics.setdefault(name, {}).setdefault( + metric_name, [] + ).append(value) + + return outer_fold_metrics + + def _score_one_inner_trial(self, model, x_folds, y_folds, metric) -> float: + """The objective of an inner search: the fold loop, recording nothing. + + The partitions of an inner trial belong to one outer fold. Rows written + from them would sit in the same table as the rows describing the model + that gets kept, indistinguishable from them and far more numerous. + """ + return self._score_folds(model, x_folds, y_folds, metric, record=False) diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index 7c33c0497..6cbb2bb43 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -10,6 +10,7 @@ "BuildModelUnit", "FitModelUnit", "FitModelOverFoldsUnit", + "FitModelOverNestedFoldsUnit", "EvaluateModelUnit", "EvaluateModelToArtifactUnit", "SaveModelUnit", @@ -234,6 +235,8 @@ def get_unit_classes(client: TestClient): ("FitModelUnit", "run_id"), ("FitModelOverFoldsUnit", "run_id"), ("FitModelOverFoldsUnit", "artifact_prefix"), + ("FitModelOverNestedFoldsUnit", "run_id"), + ("FitModelOverNestedFoldsUnit", "artifact_prefix"), ("FitModelUnit", "artifact_prefix"), ("SaveModelUnit", "artifact_prefix"), } diff --git a/tests/back/units/test_fit_model_over_folds_unit.py b/tests/back/units/test_fit_model_over_folds_unit.py new file mode 100644 index 000000000..cb9587ff2 --- /dev/null +++ b/tests/back/units/test_fit_model_over_folds_unit.py @@ -0,0 +1,310 @@ +"""Contract tests for the unit that fits a model across folds. + +Built on a hand-made ``ExecutionContext`` rather than through a job. The +end-to-end net already runs this unit inside a real run; what it cannot show is +what the unit reads, promises and refuses on its own, which is where the +composability mistakes live. +""" + +import pytest +from kink import di + +from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.fit_model_over_folds_unit import FitModelOverFoldsUnit + +#: Three folds and the trailing entry, which is not one: it holds every row the +#: folds could use and the rows reserved from them. +FOLDS = [ + {"train": "train-0", "validation": "val-0"}, + {"train": "train-1", "validation": "val-1"}, + {"train": "train-2", "validation": "val-2"}, + {"train": "pool", "test": "reserved"}, +] + + +class _RecordingModel: + """Records what was asked of it, and scores by naming what it saw.""" + + def __init__(self, run_id=7, metrics=("Accuracy",)): + self.run_id = run_id + self._metrics = list(metrics) + self.fits = [] + self.logged = [] + self.saved = [] + self.x_data = None + self.y_data = None + self._score = 0.0 + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + self.fits.append({"train": x_train, "validation": x_validation}) + + def predict(self, x_data): + return f"predictions-for-{x_data}" + + def prepare_output(self, y_data, is_fit=False): + return f"expected-from-{y_data}" + + def compute_metrics(self, split, x_data=None, y_data=None): + if not self._metrics: + return None + # A different number every time, so a mean over folds is not the same + # as any one of them. + self._score += 1.0 + return dict.fromkeys(self._metrics, self._score) + + def calculate_metrics(self, split, level, fold_index=None, **kwargs): + if not self.run_id: + return None + results = self.compute_metrics(split) + if results is None: + return None + self.logged.append((split, level, fold_index)) + return results + + def _save_metrics(self, results, split, level, **kwargs): + self.saved.append((split, level, results)) + + +class _Metric: + """Scores by naming what it was given.""" + + __name__ = "Accuracy" + + @staticmethod + def score(expected, predictions): + return 1.0 + + +class _Factory: + @staticmethod + def update_parameters(old, best): + return dict(old) + + +def _unit(**overrides): + config = { + "optimizer": {"component": "AnOptimizer", "params": {}}, + "goal_metric": "Accuracy", + "run_id": 7, + "artifact_prefix": "7", + } + config.update(overrides) + return FitModelOverFoldsUnit(**config) + + +def _context(model, folds=None): + ctx = ExecutionContext() + ctx.put("model", model) + ctx.put("x_folds", folds if folds is not None else FOLDS) + ctx.put("y_folds", folds if folds is not None else FOLDS) + ctx.put("optimizable_parameters", []) + ctx.put("factory", _Factory) + ctx.put_ref("model_parameters", {}) + return ctx + + +# --------------------------------------------------------------------------- # +# The fold loop +# --------------------------------------------------------------------------- # + + +def test_every_fold_is_fitted_and_the_kept_model_once_more(): + """Three folds and a refit, and the refit is on the pooled rows.""" + model = _RecordingModel() + + _unit()(_context(model)) + + assert [fit["train"] for fit in model.fits] == [ + "train-0", + "train-1", + "train-2", + "pool", + ] + + +def test_the_trailing_entry_is_never_scored_as_a_fold(): + """It is not a fold: it fits the model that gets kept. + + Scoring it as one would put the pooled rows -- which every fold trained on + -- into the mean that is meant to describe how the model does on rows it + has not seen. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + fold_indexes = sorted( + {index for _, level, index in model.logged if index is not None} + ) + assert fold_indexes == [0, 1, 2] + + +def test_no_fold_fit_receives_validation_data(): + """A fold is scored on the rows it held back, so the fit may not see them. + + Nothing raises when it does; the score simply comes out better than the + model deserves. Unlike the holdout sibling this is not configurable, + because there is no reading of a fold under which it would be right. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + assert all(fit["validation"] is None for fit in model.fits) + + +def test_the_fold_rows_are_written_at_the_fold_level_and_indexed_from_zero(): + """Contiguous from zero, which the repeated-cross-validation charts assume. + + They bucket folds into repetitions by integer division, so a gap moves a + fold into the wrong repetition rather than raising. + """ + model = _RecordingModel() + + _unit()(_context(model)) + + by_split = {} + for split, level, index in model.logged: + assert level is LevelEnum.FOLD + by_split.setdefault(split, []).append(index) + + assert set(by_split) == {SplitEnum.TRAIN, SplitEnum.VALIDATION} + for split, indexes in by_split.items(): + assert indexes == [0, 1, 2], split + + +def test_the_per_fold_scores_are_published_for_whoever_aggregates_them(): + """The unit does not summarise them itself. + + A summary row carries a standard deviation, and the one sanctioned write in + the domain layer has nowhere to put one -- so the numbers are handed over + and the caller persists the summary. + """ + model = _RecordingModel() + ctx = _context(model) + + _unit()(ctx) + + fold_metrics = ctx.require("fold_metrics") + assert set(fold_metrics) == {"TRAIN", "VALIDATION"} + for split_name, by_metric in fold_metrics.items(): + assert set(by_metric) == {"Accuracy"}, split_name + assert len(by_metric["Accuracy"]) == 3, split_name + + +def test_the_model_ends_pointing_at_the_partitions_it_was_last_fitted_on(): + """Whatever scores it afterwards reads the data off the instance.""" + model = _RecordingModel() + + _unit()(_context(model)) + + assert model.x_data is FOLDS[-1] + assert model.y_data is FOLDS[-1] + + +def test_a_split_with_no_metrics_configured_is_absent_rather_than_empty(): + """Present and empty would claim it was scored and produced nothing.""" + model = _RecordingModel(metrics=()) + ctx = _context(model) + + _unit()(ctx) + + assert ctx.require("fold_metrics") == {} + + +# --------------------------------------------------------------------------- # +# The objective the search measures +# --------------------------------------------------------------------------- # + + +def test_one_trial_fits_every_fold_and_returns_their_mean(): + """A trial costs k fits. That is what cross-validation buys, and its price.""" + model = _RecordingModel() + + score = _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert [fit["train"] for fit in model.fits] == ["train-0", "train-1", "train-2"] + assert score == 1.0 + + +def test_one_trial_records_one_row_per_split_holding_the_mean_of_its_folds(): + """Not one row per fold: those measure a setting, not the kept model.""" + model = _RecordingModel() + + _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert [(split, level) for split, level, _ in model.saved] == [ + (SplitEnum.TRAIN, LevelEnum.TRIAL), + (SplitEnum.VALIDATION, LevelEnum.TRIAL), + ] + # Three folds scoring 1, 3, 5 on TRAIN (the double counts every call) means + # the row is their mean and not any one of them. + train_row = model.saved[0][2] + assert train_row["Accuracy"] == 3.0 + + +def test_a_trial_writes_no_rows_when_there_is_no_run_to_write_against(): + """``_save_metrics`` does not guard this for itself. + + A caller with no run -- a pipeline -- would write rows against a foreign + key pointing at nothing, and they insert without complaint because nothing + enforces it. That is the failure a metrics unit already exists to refuse. + """ + model = _RecordingModel(run_id=None) + + _unit()._score_one_trial(model, FOLDS, FOLDS, _Metric) + + assert model.saved == [] + assert model.fits, "the folds are still fitted; only the recording is refused" + + +# --------------------------------------------------------------------------- # +# The contract itself +# --------------------------------------------------------------------------- # + + +def test_the_unit_refuses_to_run_before_the_model_was_built(): + """``__call__`` checks REQUIRES, so the failure names the missing key.""" + with pytest.raises(UnitContractError): + _unit()(ExecutionContext()) + + +def test_validate_raises_when_called_before_the_model_was_built(): + """A missing key is a call-order mistake, not "nothing to optimize". + + The two are different: the key present and empty is a model that declares + no optimizable parameters, and that one legitimately skips the checks. + """ + with pytest.raises(UnitContractError, match="'optimizable_parameters'"): + _unit().validate(ExecutionContext()) + + +def test_two_units_in_one_context_do_not_share_their_optimizer(): + """Instance state stays on the instance, which is what lets a graph hold two. + + A context-global cache would give the second unit the first one's + optimizer, and the second search would silently run the first one's. + """ + + class _Optimizer: + def __init__(self, **params): + pass + + registry = { + "AnOptimizer": {"class": _Optimizer}, + "Accuracy": {"class": _Metric, "metadata": {"maximize": True}}, + } + di["component_registry"] = registry + try: + ctx = ExecutionContext() + ctx.put("optimizable_parameters", ["lr"]) + + first, second = _unit(), _unit() + first.validate(ctx) + second.validate(ctx) + + assert first._optimizer is not second._optimizer + assert not ctx.has("optimizer") + finally: + del di["component_registry"] From d476cf0a0ff3f4ab659c6b2cf881ff395247b2dd Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 21:10:08 -0300 Subject: [PATCH 26/28] Leave the evaluation strategies declaring, and nothing else They ran the training: execute took the run row and the database session and did the fitting, the search, the scoring, the aggregation and the persistence behind one method. Every piece of that is now a unit, and nothing has called execute since the fold paths moved -- the job reads SCORED_SPLITS and KIND off the class it resolves and never touches it otherwise. So the code goes. base_evaluation_strategy, cv and holdout drop from 881 lines to 153, and what is left is what was underneath all along: how a run is carved, and which partitions it records a score for. They stay registered. That is not deference to dead code -- the frontend reads these classes in four places, and only one is about metrics. The session wizard lists them so the user can choose one, and ModelSession.evaluation_strategy is NOT NULL, so without that listing a session cannot be created at all. It starts on the first one whose kind is holdout. `kind` decides the shape of the splits payload and which controls are shown. Only `scored_splits` is about the charts. Removing the classes would not have cost two screens; it would have cost the way sessions are made. There is precedent for a class here that declares and does not execute: BaseSplitter.PARTITIONING and explainable_partitions are read exactly this way, by the backend and by the frontend, and nothing calls them to do work. Five of the forecasting tests exercised behaviour rather than declarations -- the final fit, and which partitions a trial scores. That behaviour moved rather than disappeared, so they are pointed at the units that carry it out now. They stay in the same file, next to the declarations, because that is the pair that has to stay consistent: a strategy that says it does not score the training partition, and a fit that then does not. 1014 passed across units, dag, spike, api and evaluation. Co-Authored-By: Claude Opus 5 (1M context) --- .../evaluation/base_evaluation_strategy.py | 257 ++-------- DashAI/back/evaluation/cv.py | 441 +----------------- DashAI/back/evaluation/holdout.py | 172 +------ .../evaluation/test_forecasting_strategies.py | 68 ++- 4 files changed, 122 insertions(+), 816 deletions(-) diff --git a/DashAI/back/evaluation/base_evaluation_strategy.py b/DashAI/back/evaluation/base_evaluation_strategy.py index f100bad2d..6b07240b9 100644 --- a/DashAI/back/evaluation/base_evaluation_strategy.py +++ b/DashAI/back/evaluation/base_evaluation_strategy.py @@ -1,31 +1,56 @@ -import os -import pickle -from abc import ABCMeta, abstractmethod -from typing import Callable, Final, List, Optional +"""What a run records, declared per strategy. -from kink import di +These classes used to run the training too: ``execute`` took the run row and the +database session and did the fitting, the search, the scoring, the aggregation +and the persistence behind one method. That work is the units' now -- see +``FitModelUnit``, ``FitModelOverFoldsUnit`` and their siblings -- and what is +left here is what was underneath it all along: a declaration of how a run is +carved and which partitions it records a score for. + +**They stay registered even though they no longer do anything.** The frontend +reads them in four places, and only one is about metrics: + +- the session wizard lists them so the user can choose one, and + ``ModelSession.evaluation_strategy`` is NOT NULL, so without that listing a + session cannot be created at all; +- it starts on the first one whose ``kind`` is holdout; +- ``kind`` decides the shape of the splits payload and which controls are shown; +- ``scored_splits`` tells the metric charts which partitions exist to plot. + +There is precedent for a class in this codebase that declares and does not +execute: ``BaseSplitter.PARTITIONING`` and ``explainable_partitions`` are read +the same way, by the backend and by the frontend, and nothing calls them to do +work. +""" + +from typing import Final -from DashAI.back.core.artifacts import normalize_artifacts from DashAI.back.core.enums.metrics import SplitEnum -from DashAI.back.dependencies.database.models import Run -from DashAI.back.models.base_model import BaseModel -from DashAI.back.models.model_factory import ModelFactory -from DashAI.back.optimizers.base_optimizer import BaseOptimizer -class BaseEvaluationStrategy(metaclass=ABCMeta): - """Abstract base class defining the interface for model evaluation strategies. +class BaseEvaluationStrategy: + """How a run is carved, and what it records. - Concrete implementations (e.g., CrossValidationEvaluationStrategy, - HoldoutEvaluationStrategy) inherit from this class and provide specific - strategies for model evaluation. + Subclasses declare; none of them execute. """ TYPE: Final[str] = "EvaluationStrategy" + #: Whether this strategy splits the dataset once or into folds. It decides + #: which unit prepares the data, because the two publish different shapes, + #: and which controls the session wizard offers. KIND: str = "holdout" + + #: The partitions a run records a score for. Not the same as which + #: partitions have metrics configured: a forecaster has training metrics + #: and still must not be judged on the dates it was fitted on, because an + #: in-sample fit statistic is not comparable with a forecast. A screen that + #: offers one control per partition reads this instead of assuming all + #: three exist. SCORED_SPLITS: tuple = (SplitEnum.TRAIN, SplitEnum.VALIDATION, SplitEnum.TEST) + #: The partition the kept model was fitted on, which is what decides which + #: partitions of a finished run can still be predicted. FINAL_FIT_PARTITIONS: tuple = ("train",) @classmethod @@ -37,209 +62,9 @@ def get_metadata(cls) -> dict: dict Mapping with ``kind``, which says whether this strategy splits the dataset once or into folds, and ``scored_splits``, the partitions - it writes metrics for. A screen that offers one control per - partition reads the latter instead of assuming all three exist: - a forecasting strategy scores no training partition, so asking it - for train metrics finds nothing. + it writes metrics for. """ return { "kind": cls.KIND, "scored_splits": [split.value for split in cls.SCORED_SPLITS], } - - def __init__( - self, - factory: ModelFactory, - optimizer: BaseOptimizer, - goal_metric, - **kwargs, - ): - """Initialize the evaluation strategy with model and optimization configuration. - - Parameters - ---------- - factory : ModelFactory - Factory owning the model to be trained/evaluated and the - hyperparameters that are eligible for optimization. - optimizer : BaseOptimizer - The hyperparameter optimizer instance. Can be None if no HPO is needed. - goal_metric : dict (obtained from Metric component registry) - The target metric to optimize during hyperparameter search. - **kwargs - Additional keyword arguments passed from subclasses (ignored). - """ - self.factory: ModelFactory = factory - self.model: BaseModel = factory.model - self.run_optimizable_parameters = factory.optimizable_parameters - self.optimizer: BaseOptimizer = optimizer - self.goal_metric = goal_metric - self._progress_reporter: Optional[ - Callable[[Optional[float], Optional[str]], None] - ] = None - - def set_progress_reporter( - self, - progress_reporter: Optional[Callable[[Optional[float], Optional[str]], None]], - ) -> None: - """Register a callback that will receive progress updates.""" - self._progress_reporter = progress_reporter - - def _report_progress( - self, fraction: Optional[float], message: Optional[str] = None - ): - """Emit progress updates when a reporter has been registered.""" - if self._progress_reporter is not None: - self._progress_reporter(fraction, message) - - @abstractmethod - def execute(self, x, y, run: Run, db): - """Execute the evaluation strategy on the provided data. - - This is the main entry point for the evaluation process. Subclasses implement - strategy-specific logic for: - - Model training across folds/splits - - Metric computation and persistence - - HPO execution and result handling - - Parameters - ---------- - x : DatasetDict or list of DastasetDict - Input features. Structure depends on the evaluation strategy: - - For holdout: DatasetDict with train/validation/test splits - - For CV: List of DatasetDicts, one per fold with train/test splits - y : dict or list - Target labels. Same structure as x. - run : Run - Database model representing the current experiment run. - db : Session - SQLAlchemy database session for persisting results. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model after evaluation - - plot_paths : list[str] - Paths to generated HPO visualization files - """ - raise NotImplementedError("Subclasses must implement this method") - - @abstractmethod - def evaluate(self, model: BaseModel, x, y, metric): - """Evaluate the model on the given data and return the score. - - This method is called during hyperparameter optimization to compute - the objective function value for a given set of hyperparameters. - Different strategies may compute metrics differently (e.g., across CV folds - or on a validation split). - - Parameters - ---------- - model : BaseModel - The model instance to evaluate. - x : DatasetDict or list of DastasetDict - Input features for evaluation (structure depends on strategy). - y : DatasetDict or list of DastasetDict - Target labels for evaluation (structure depends on strategy). - metric : Metric - The metric instance to compute. - - Returns - ------- - float - The computed metric value used as the optimization objective. - """ - raise NotImplementedError("Subclasses must implement this method") - - def _do_hpo(self, model: BaseModel, x, y, run: Run, db): - """Execute hyperparameter optimization using the configured optimizer. - - The optimizer uses the self.evaluate method as the objective function, - allowing each strategy to define its own evaluation logic. - - Parameters - ---------- - model : BaseModel - The model instance to optimize. - x : DatasetDict or list of DatasetDict - Training input features (structure varies by strategy). - y : DatasetDict or list of DatasetDict - Training target labels (structure varies by strategy). - run : Run - Database run instance to update with optimized parameters. - db : Session - SQLAlchemy database session for transactions. - - Returns - ------- - BaseModel - The model with the best hyperparameters found during optimization. - """ - from sqlalchemy.orm.attributes import flag_modified - - # Execute hyperparameter optimization and get best model with parameters - self.optimizer.optimize( - model, - x, - y, - self.run_optimizable_parameters, - self.goal_metric, - strategy=self.evaluate, - ) - model = self.optimizer.get_model() - best_params = self.optimizer.get_best_params() - - # Update the run's parameters with the optimized hyperparameters - old_parameters = run.parameters.copy() - updated_parameters = self.factory.update_parameters(old_parameters, best_params) - - # Persist the updated parameters to the database - run.parameters = updated_parameters - flag_modified(run, "parameters") - db.commit() - - return model - - def _generate_hpo_plots(self, run: Run) -> List[str]: - """Generate and pickle the hyperparameter optimization plots to disk. - - Shared by every evaluation strategy that runs HPO, so the plot - generation logic only needs to be maintained in one place. - - Parameters - ---------- - run : Run - The run the plots belong to (used for the plot filenames). - - Returns - ------- - list[str] - Paths to the pickled plot files, in the order produced by the - optimizer. - """ - config = di["config"] - plot_paths: List[str] = [] - - # Retrieve optimization trial data from the optimizer - trials = self.optimizer.get_trials_values() - - # Generate plot visualizations from the trial data - # Plots typically show parameter importance, optimization history, etc. - plot_filenames, plots = self.optimizer.create_plots( - trials, - run.id, - n_params=len(self.run_optimizable_parameters), - goal_metric=self.goal_metric, - ) - - # Convert plots to serializable format (handles special objects, arrays, etc.) - normalized_plots = normalize_artifacts(plots) - - # Serialize and persist each plot to disk - for filename, plot in zip(plot_filenames, normalized_plots, strict=False): - plot_path = os.path.join(config["RUNS_PATH"], filename) - # Serialize the plot object using pickle and write to disk - with open(plot_path, "wb") as file: - pickle.dump(plot, file) - plot_paths.append(plot_path) - - return plot_paths diff --git a/DashAI/back/evaluation/cv.py b/DashAI/back/evaluation/cv.py index 54b771136..39e0e8b7e 100644 --- a/DashAI/back/evaluation/cv.py +++ b/DashAI/back/evaluation/cv.py @@ -1,434 +1,31 @@ -from functools import partial +"""Cross-validation: several train and validation pairs, scored per fold.""" -import numpy as np -from kink import di - -from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum -from DashAI.back.dependencies.database.models import Metric, Run from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy -from DashAI.back.splitters.base_splitter import BaseSplitter class FoldEvaluationStrategy(BaseEvaluationStrategy): - """Evaluation strategy implementing k-fold cross-validation with optional - nested CV and HPO. - - This strategy partitions the dataset into k folds and performs k rounds of - training and evaluation. Each fold is split into a train and a validation - partition, so the scores obtained by resampling are validation estimates. - When the session reserved rows, the final model is fitted on everything the - folds could use and scored once against those reserved rows. - - The strategy handles metric aggregation at multiple levels: - - FOLD level: Individual metrics from each fold - - TRIAL level: Metrics during HPO trials - - LAST/LAST_OUTER: Aggregated metrics (mean and std) for simple/nested CV + """Carve the dataset into folds and score each of them. + + Each fold is split into a train and a validation partition, so the scores + obtained by resampling are validation estimates. When the session reserved + rows, the kept model is fitted on everything the folds could use and scored + once against those reserved rows -- the only data no fold and no trial ever + saw. + + A run carved this way is prepared by ``PrepareAndFoldUnit`` and fitted by + ``FitModelOverFoldsUnit``, or by ``FitModelOverNestedFoldsUnit`` when the + run also asks for a search inside each fold. + + The levels a fold run records, which is why the enum has more of them than + a holdout run needs: ``FOLD`` per fold, aggregated to ``LAST`` with a + standard deviation; ``TRIAL`` once per trial of a search, holding the mean + over that trial's folds; and, for a nested run, ``OUTER_FOLD`` aggregated + to ``LAST_OUTER``, kept apart because it answers a different question -- + how the procedure does rather than how this model does. """ KIND: str = "cv" - def execute(self, x, y, run: Run, db): - """Execute k-fold cross-validation with optional nested CV and HPO. - - Trains and evaluates a model using k-fold cross-validation. Optionally performs - hyperparameter optimization and nested CV to prevent overfitting. Aggregates - metrics across folds and returns the trained model. - - Parameters - ---------- - x : list of DatasetDict - List of fold DatasetDict each containing: - {"train": X_train, "validation": X_validation} - The last element is not a fold: it holds every row the folds could - use as {"train": X_pool, "test": X_test}, where the test partition - is empty when the session reserved nothing. - y : list of DatasetDict - List of fold label DatasetDicts with same structure as x. - run : Run - Database run instance containing configuration (nested CV settings, etc.). - db : Session - SQLAlchemy database session for persisting metrics and parameters. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model - - plot_paths : list[str] - Paths to HPO visualization plot files - """ - plot_paths = [] - model = self.model - - # STEP 1: Hyperparameter Optimization (if enabled) - if self.optimizer and self.run_optimizable_parameters: - # Initialize nested CV if required - if run.nested: - try: - registry = di["component_registry"] - inner_splits = run.nested - splitter_name = inner_splits.get("splitter_name", None) - - # Create inner splitter for nested CV fold generation - self.inner_splitter: BaseSplitter = registry[splitter_name][ - "class" - ](inner_splits) - except Exception as e: - raise ValueError( - f"Error configuring inner splitter for nested CV: {e}" - ) from e - - # Execute nested cross-validation for HPO - self._report_progress(0.1, "Nested cross-validation") - self._nested_cv(run.id, model, x, y, db) - - # Perform hyperparameter optimization and update run.parameters - self._report_progress(0.25, "Hyperparameter optimization") - model = self._do_hpo(model, x, y, run, db) - - # Generate and serialize HPO visualization plots - plot_paths = self._generate_hpo_plots(run) - - total_folds = len(x) - 1 - - # STEP 2: Main k-fold Cross-Validation Loop - # Note: Last fold (index len(x)-1) is reserved for final training, - # not CV evaluation - for i in range(total_folds): - self._report_progress( - 0.4 + ((i + 1) / total_folds) * 0.4, - f"Evaluating fold {i + 1}/{total_folds}", - ) - x_fold = x[i] - y_fold = y[i] - - # Set model's internal references to current fold data - model.x_data = x_fold - model.y_data = y_fold - - # Train model on fold's training partition - model.train(x_fold["train"], y_fold["train"]) - - # Compute and store metrics for this fold - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model.calculate_metrics( - split=SplitEnum.TRAIN, level=LevelEnum.FOLD, fold_index=i - ) - model.calculate_metrics( - split=SplitEnum.VALIDATION, level=LevelEnum.FOLD, fold_index=i - ) - - # STEP 3: Aggregate metrics across all folds - # Compute mean and std of fold metrics and store as LAST level metrics - self._aggregate_fold_metrics( - run_id=run.id, - db=db, - level_to_agg=LevelEnum.FOLD, - level_to_save=LevelEnum.LAST, - ) - - # STEP 4: Final model training on every row the folds could use, which - # excludes the rows reserved when the session asked for them. - self._report_progress(0.85, "Training final model") - model.x_data = x[-1] - model.y_data = y[-1] - model.train(x[-1]["train"], y[-1]["train"]) - - # STEP 5: Score the final model once on the reserved rows. No fold and - # no hyperparameter trial ever saw them, so this is the only estimate - # that stays honest after a model is picked out of the comparison - # table. It is a single measurement, hence no standard deviation. - if len(x[-1]["test"]) > 0: - self._report_progress(0.90, "Evaluating on the reserved rows") - model.calculate_metrics(split=SplitEnum.TEST, level=LevelEnum.LAST) - - return model, plot_paths - - def evaluate(self, model, input_dataset, output_dataset, metric, **kwargs): - """Evaluate model using k-fold cross-validation (used as HPO objective - function). - - This method implements cross-validation evaluation for hyperparameter - optimization. It trains on each fold's train partition, scores the fold's - validation partition and averages the results. When used in nested CV, it - evaluates on inner folds within a specific outer fold. - - Parameters - ---------- - model : BaseModel - The model instance to evaluate (with specific hyperparameters). - input_dataset : list of DatasetDict - List of fold data {"train": X_train, "validation": X_validation}. - output_dataset : list of DatasetDict - List of fold labels {"train": y_train, "validation": y_validation}. - metric : Metric - The metric function to optimize. - **kwargs - Additional arguments including: - - fold_index : int or None - Inner outer fold index in nested CV (None for simple CV) - - Returns - ------- - float - Mean metric value across all k folds (objective value for HPO). - - Note: When fold_index is provided (nested CV inner loop), - intermediate metrics are NOT being saved (only outer loop metrics are saved). - """ - # Extract context: fold_index indicates if we're in nested CV inner loop - # None means simple CV; an integer means nested CV on that outer fold - fold_index = kwargs.get("fold_index") - - # List to collect the goal metric value from each fold - folds_results = [] - - # Dictionaries to accumulate all metrics across folds for averaging - train_results = {} - validation_results = {} - - # Cross-validation loop - # Iterate through k-1 folds (last fold is the complete dataset) - # This loop represents either: - # - Main CV loop (if fold_index is None) - # - Inner CV loop within outer fold i (if fold_index is set) - for i in range(len(input_dataset) - 1): - x_fold = input_dataset[i] - y_fold = output_dataset[i] - - # Set model's internal data references for this fold - model.x_data = x_fold - model.y_data = y_fold - - # Train model on this fold's training partition - model.train(x_fold["train"], y_fold["train"]) - - # Compute metrics on both training and validation sets. - # - # ``compute_metrics`` returns None when there was nothing to score - # at all -- no metrics configured, or a partition with no rows -- - # which is not the same statement as "scored, and the result was - # empty". A fold with nothing to score cannot contribute a number - # to the mean below, and averaging over the folds that did have - # data would report a score for a search the user cannot reproduce. - train_scores = ( - model.compute_metrics(split=SplitEnum.TRAIN) - if SplitEnum.TRAIN in self.SCORED_SPLITS - else {} - ) - validation_scores = model.compute_metrics(split=SplitEnum.VALIDATION) - if validation_scores is None: - raise ValueError( - f"Fold {i} has no validation data to score, so the " - "cross-validation objective cannot be computed. Check the " - "number of folds against the size of the dataset." - ) - - # Collect the goal metric value from this fold - folds_results.append(validation_scores[metric.__name__]) - - # Accumulate all metrics only if NOT in nested CV inner loop - if fold_index is None: - for results, scores in [ - (train_results, train_scores or {}), - (validation_results, validation_scores), - ]: - for metric_name, value in scores.items(): - if metric_name not in results: - results[metric_name] = [] - results[metric_name].append(value) - - # Save intermediate metrics (simple CV only) - if fold_index is None: - # Compute average metrics across all folds - averaged_train_results = { - metric: np.mean(values) for metric, values in train_results.items() - } - averaged_validation_results = { - metric: np.mean(values) for metric, values in validation_results.items() - } - - # Persist averaged metrics as TRIAL level (intermediate HPO result) - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model._save_metrics( - results=averaged_train_results, - split=SplitEnum.TRAIN, - level=LevelEnum.TRIAL, - ) - model._save_metrics( - results=averaged_validation_results, - split=SplitEnum.VALIDATION, - level=LevelEnum.TRIAL, - ) - - # Return the mean of the goal metric across folds - # This is the objective value used by the optimizer - return np.mean(folds_results) - - def _nested_cv(self, run_id, model, input_dataset, output_dataset, db): - """Execute nested cross-validation with inner HPO loop. - - Nested CV implements a two-level validation scheme to prevent overfitting during - hyperparameter optimization and give an unbiased estimate of model performance. - - Outer loop: Standard k-fold CV for unbiased final performance estimation - - Inner loop: Separate CV fold within each outer fold for HPO - - This prevents "information leakage" where the data used to score a fold also - influences hyperparameter selection, which would overestimate true - generalization performance. - - Parameters - ---------- - run_id : int - The database run ID for metric storage. - model : BaseModel - The model instance to optimize and evaluate. - input_dataset : list of DatasetDict - List of outer fold data - {"train": X_train, "validation": X_validation}. - output_dataset : list of DatasetDict - List of outer fold labels - {"train": y_train, "validation": y_validation}. - db : Session - SQLAlchemy database session for persisting aggregated metrics. - """ - # Nested CV Outer Loop - # For each outer fold, optimize hyperparameters on inner folds - for i in range(len(input_dataset) - 1): - x_outer = input_dataset[i] - y_outer = output_dataset[i] - - # Create inner folds from outer fold's training data - # This ensures HPO validation data is independent of outer test fold - inner_x, inner_y, _ = self.inner_splitter.split( - x_outer["train"], y_outer["train"] - ) - - # Create evaluation strategy for this outer fold - # Passes fold_index so evaluate() knows it's in nested CV context - strategy_with_context = partial(self.evaluate, fold_index=i) - - # INNER LOOP: Optimize hyperparameters using inner CV - # The optimizer will iteratively: - # - Generate hyperparameter candidates - # - Train models on inner fold combinations - # - Evaluate using inner CV (calls strategy_with_context) - # - Select hyperparameters with best inner CV performance - self.optimizer.optimize( - model, - inner_x, - inner_y, - self.run_optimizable_parameters, - self.goal_metric, - strategy=strategy_with_context, - ) - outer_model = self.optimizer.get_model() - - # Set model's data references for outer fold evaluation - outer_model.x_data = x_outer - outer_model.y_data = y_outer - - # Train model on outer fold's training data with optimized hyperparameters - outer_model.train(x_outer["train"], y_outer["train"]) - - # Evaluate on outer fold's test data (this is OUTER_FOLD level metric) - outer_model.calculate_metrics( - split=SplitEnum.VALIDATION, level=LevelEnum.OUTER_FOLD, fold_index=i - ) - if SplitEnum.TRAIN in self.SCORED_SPLITS: - outer_model.calculate_metrics( - split=SplitEnum.TRAIN, level=LevelEnum.OUTER_FOLD, fold_index=i - ) - - # Aggregate outer fold metrics - # Compute mean and std of OUTER_FOLD metrics and store as LAST_OUTER level - self._aggregate_fold_metrics( - run_id=run_id, - db=db, - level_to_agg=LevelEnum.OUTER_FOLD, - level_to_save=LevelEnum.LAST_OUTER, - ) - - def _aggregate_fold_metrics( - self, - run_id: int, - db, - level_to_agg=LevelEnum.FOLD, - level_to_save=LevelEnum.LAST, - ): - """Aggregate and average fold metrics across cross-validation folds. - - This method computes the mean and standard deviation of metrics collected - at the fold level and stores the aggregated results at a higher level. - This is used to provide summary statistics for model performance. - - Typical usage patterns: - - Aggregate FOLD metrics -> store as LAST (simple CV summary) - - Aggregate OUTER_FOLD metrics -> store as LAST_OUTER (nested CV summary) - - Parameters - ---------- - run_id : int - The database run ID to aggregate metrics for. - db : Session - SQLAlchemy database session for querying and persisting metrics. - level_to_agg : LevelEnum, optional - The source metric level to aggregate. Default: FOLD. - level_to_save : LevelEnum, optional - The destination level for aggregated metrics. Default: LAST. - """ - # Query all metrics with level=level_to_agg for this run - fold_metrics = ( - db.query(Metric) - .filter(Metric.run_id == run_id, Metric.level == level_to_agg) - .all() - ) - - # If no metrics found, nothing to aggregate - if not fold_metrics: - return - - # Group metrics by (split, name) for aggregation - metrics_by_split_name = {} - for metric in fold_metrics: - key = (metric.split, metric.name) - if key not in metrics_by_split_name: - metrics_by_split_name[key] = [] - metrics_by_split_name[key].append(metric.value) - - # Aggregate and persist metrics - for (split, name), values in metrics_by_split_name.items(): - # Compute aggregation statistics - avg_value = np.mean(values) - std_value = np.std(values) if len(values) > 1 else 0.0 - - # Check if aggregated metric already exists for this split/name/run - existing = ( - db.query(Metric) - .filter_by(run_id=run_id, split=split, level=level_to_save, name=name) - .first() - ) - - if existing: - # Update existing metric with aggregated values - existing.value = avg_value - existing.std_value = std_value - else: - # Create new aggregated metric - db.add( - Metric( - run_id=run_id, - split=split, - level=level_to_save, - name=name, - value=avg_value, - std_value=std_value, - step=0, - ) - ) - - # Persist aggregated metrics to database - db.commit() - class CrossValidationEvaluationStrategy(FoldEvaluationStrategy): """Score a model across folds, recording train and validation for each. diff --git a/DashAI/back/evaluation/holdout.py b/DashAI/back/evaluation/holdout.py index 6fb602ccc..52b9b7e4c 100644 --- a/DashAI/back/evaluation/holdout.py +++ b/DashAI/back/evaluation/holdout.py @@ -1,171 +1,19 @@ -from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum -from DashAI.back.dependencies.database.models import Metric, Run +"""Holdout evaluation: one split into three partitions, scored once.""" + from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy class SinglePartitionEvaluationStrategy(BaseEvaluationStrategy): - """Evaluation strategy implementing holdout (train/validation/test split) - validation. - - This strategy divides the dataset into three mutually exclusive partitions: - training, validation, and test. The training set is used for model training, - the validation set for HPO, and the test set for final evaluation. + """Split once into train, validation and test. - The strategy handles metric aggregation at multiple levels: - - TRIAL level: Metrics during HPO trials on validation set - - LAST level: Final metrics computed on all three partitions after training + The training set fits the model, the validation set is what a + hyperparameter search is measured on, and the test set is scored once at + the end. A run carved this way is prepared by ``PrepareAndSplitUnit`` and + fitted by ``FitModelUnit``. """ KIND: str = "holdout" - def execute(self, x, y, run: Run, db): - """Execute holdout validation: train on training set, optimize with validation, - evaluate on test. - - Trains a model on the training partition and optionally performs hyperparameter - optimization using the validation set. Finally evaluates the trained model on - all three partitions (train, validation, test) and returns the model with plots. - - Parameters - ---------- - x : DatasetDict - DatasetDict with data partitions: - {"train": X_train, "validation": X_val, "test": X_test} - y : DatasetDict - DatasetDict with label partitions: - {"train": y_train, "validation": y_val, "test": y_test} - run : Run - Database run instance for storing results and configuration. - db : Session - SQLAlchemy database session for persisting metrics. - - Returns - ------- - tuple - (trained_model, plot_paths) where: - - trained_model : BaseModel - The trained model - - plot_paths : list[str] - Paths to HPO visualization plot files - """ - plot_paths = [] - model = self.model - - # set the data used for model training and evaluation - model.x_data = x - model.y_data = y - - # Execute HPO if optimizer and there are parameters to optimize - if self.optimizer and self.run_optimizable_parameters: - self._report_progress(0.2, "Hyperparameter optimization") - model = self._do_hpo(model, x, y, run, db) - plot_paths = self._generate_hpo_plots(run) - - # Train the model with the provided data and return it - self._report_progress(0.5, "Training") - self._fit_final_model(model, x, y) - - # Calculate metrics at the end of training if not done already - self._report_progress(0.85, "Computing metrics") - for split in self.SCORED_SPLITS: - self._calculate_metrics_if_missing(model, run, db, split) - - return model, plot_paths - - def _fit_final_model(self, model, x, y): - """Fit the model that gets kept. - - Separate from the trial fits so a strategy can differ on what the - kept model is allowed to learn from, which is the one thing - forecasting needs to change here. - - Parameters - ---------- - model : BaseModel - The model to fit. - x : DatasetDict - Input partitions. - y : DatasetDict - Target partitions. - """ - model.train(x["train"], y["train"], x["validation"], y["validation"]) - - def _calculate_metrics_if_missing(self, model, run: Run, db, split: SplitEnum): - """Compute and persist LAST-level metrics for a split, unless already saved. - - Parameters - ---------- - model : BaseModel - The trained model to compute metrics on. - run : Run - Database run instance the metrics belong to. - db : Session - SQLAlchemy database session used to check for existing metrics. - split : SplitEnum - The data split to compute metrics for. - """ - existing_metric = ( - db.query(Metric) - .filter_by(run_id=run.id, split=split, level=LevelEnum.LAST) - .first() - ) - if not existing_metric: - model.calculate_metrics(split=split, level=LevelEnum.LAST) - - def evaluate(self, model, input_dataset, output_dataset, metric): - """Evaluate model on validation set during HPO trials. - - Trains the model on the training set and computes the metric on the validation - set. Used as the objective function during hyperparameter optimization. - - Parameters - ---------- - model : BaseModel - The model instance to evaluate with specific hyperparameters. - input_dataset : DatasetDict - DatasetDict with data partitions - {"train": X_train, "validation": X_val, "test": X_test}. - output_dataset : DatasetDict - DatasetDict with label partitions - {"train": y_train, "validation": y_val, "test": y_test}. - metric : Metric - The metric function to compute on predictions. - - Returns - ------- - float - The metric score value for this hyperparameter combination. - """ - # Validation data is passed on purpose. Without it the epoch loops - # skip `calculate_metrics(split=VALIDATION, level=EPOCH)` entirely — - # they guard it behind `if x_validation is not None` — so during - # optimization the per-epoch validation score was never computed. - # - # That is the deeper reason pruning could not work here: the number a - # pruner needs to decide did not exist, independently of whether the - # pruner itself was an instance or a string. - model.train( - input_dataset["train"], - output_dataset["train"], - input_dataset["validation"], - output_dataset["validation"], - ) - - # Evaluate the model on the validation set - y_pred = model.predict(input_dataset["validation"]) - - output_dataset_transformed = model.prepare_output( - output_dataset["validation"], is_fit=False - ) - - # Calculate metric for train and validation data each trial. - if SplitEnum.TRAIN in self.SCORED_SPLITS: - model.calculate_metrics(split=SplitEnum.TRAIN, level=LevelEnum.TRIAL) - model.calculate_metrics(split=SplitEnum.VALIDATION, level=LevelEnum.TRIAL) - - # Compute the objective metric score on the validation set - score = metric.score(output_dataset_transformed, y_pred) - - return score - class HoldoutEvaluationStrategy(SinglePartitionEvaluationStrategy): """Split once into train, validation and test, and score all three. @@ -176,8 +24,10 @@ class HoldoutEvaluationStrategy(SinglePartitionEvaluationStrategy): not belong in the same results table as one. ``ForecastingHoldoutEvaluationStrategy`` records validation and test only. - The final fit is the same in both: the kept model is fitted on the training - partition alone, so it is the model the recorded metrics describe. + The kept model is fitted on the training partition alone in both, so it is + the model the recorded metrics describe. The validation partition is handed + to the fit, which models use to watch training and stop early -- being + given it to watch is not the same as being fitted on it. """ COMPATIBLE_COMPONENTS = [ diff --git a/tests/back/evaluation/test_forecasting_strategies.py b/tests/back/evaluation/test_forecasting_strategies.py index a2f8b3c17..abee9c8f7 100644 --- a/tests/back/evaluation/test_forecasting_strategies.py +++ b/tests/back/evaluation/test_forecasting_strategies.py @@ -127,17 +127,40 @@ def test_the_ordinary_holdout_strategy_still_scores_all_three(): ) +def _fit_unit(**overrides): + """A FitModelUnit configured for an ordinary run. + + The strategies declare how a run is evaluated; this unit is what carries it + out. The tests below live next to the declarations because that is the pair + that has to stay consistent -- a strategy that says it does not score the + training partition, and a fit that then does not. + """ + from DashAI.back.units.fit_model_unit import FitModelUnit + + config = { + "optimizer": {"component": "", "params": {}}, + "goal_metric": "", + "run_id": None, + "artifact_prefix": None, + } + config.update(overrides) + return FitModelUnit(**config) + + # --- the final fit ----------------------------------------------------------- def test_the_kept_model_stops_at_the_end_of_training(): + """Handing the validation partition to the fit is not fitting on it. + + A model is given it to watch and stop early. A forecaster must not be + advanced through it: the kept model has to end where the training rows end, + or the validation metrics describe a model that already saw them. + """ xs, ys, _ = _split() - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) model = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(model, xs, ys) + _fit_unit()._fit_kept_model(model, xs, ys) last_train_date = pd.to_datetime(xs["train"].to_pandas().iloc[:, 0]).max() assert model._last_train_date == last_train_date @@ -147,15 +170,12 @@ def test_the_kept_model_is_the_one_the_validation_metrics_describe(): from DashAI.back.metrics.regression.mae import MAE xs, ys, _ = _split() - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) scored = ExponentialSmoothing(seasonal="add", season_length=12) scored.train(xs["train"], ys["train"]) kept = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(kept, xs, ys) + _fit_unit()._fit_kept_model(kept, xs, ys) assert MAE.score(ys["validation"], kept.predict(xs["validation"])) == MAE.score( ys["validation"], scored.predict(xs["validation"]) @@ -163,16 +183,18 @@ def test_the_kept_model_is_the_one_the_validation_metrics_describe(): def test_a_session_without_validation_rows_still_fits(): + """Present and empty, which is not the same as absent. + + The fit is handed a validation partition with no rows in it rather than + none at all, and has to cope: a session may reserve nothing for it. + """ xs, ys, _ = _split() empty = xs["validation"].select(range(0)) xs = {**xs, "validation": empty} ys = {**ys, "validation": ys["validation"].select(range(0))} - strategy = ForecastingHoldoutEvaluationStrategy.__new__( - ForecastingHoldoutEvaluationStrategy - ) model = ExponentialSmoothing(seasonal="add", season_length=12) - strategy._fit_final_model(model, xs, ys) + _fit_unit()._fit_kept_model(model, xs, ys) last_train_date = pd.to_datetime(xs["train"].to_pandas().iloc[:, 0]).max() assert model._last_train_date == last_train_date @@ -202,25 +224,37 @@ def calculate_metrics(self, split, level=None, **kwargs): self.scored.append(split) -def _evaluate_with(strategy_class): +def _trial_of(strategy_class): + """Run one trial configured the way this strategy's runs are configured. + + Which partitions a trial records is no longer the strategy's own code: the + job reads SCORED_SPLITS off the class and hands it to the fitting unit, + minus the test partition, which a trial may never score. This does the same + thing so the declaration stays connected to what it produces. + """ from DashAI.back.metrics.regression.mae import MAE xs, ys, _ = _split() - strategy = strategy_class.__new__(strategy_class) model = _RecordingModel() - strategy.evaluate(model, xs, ys, MAE) + trial_splits = [ + split.name + for split in strategy_class.SCORED_SPLITS + if split is not SplitEnum.TEST + ] + + _fit_unit(trial_splits=trial_splits)._score_one_trial(model, xs, ys, MAE) return model.scored def test_a_forecasting_trial_never_scores_the_training_partition(): - scored = _evaluate_with(ForecastingHoldoutEvaluationStrategy) + scored = _trial_of(ForecastingHoldoutEvaluationStrategy) assert SplitEnum.TRAIN not in scored assert SplitEnum.VALIDATION in scored def test_an_ordinary_trial_still_scores_both(): - scored = _evaluate_with(HoldoutEvaluationStrategy) + scored = _trial_of(HoldoutEvaluationStrategy) assert SplitEnum.TRAIN in scored assert SplitEnum.VALIDATION in scored From f30f3bcbc61e704173ff1ba40b08c625d1d5b5da Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 21:37:28 -0300 Subject: [PATCH 27/28] Build the pruning tests' objective from the unit that fits Two optimizer test files built the objective they measure by reaching for HoldoutEvaluationStrategy.evaluate. That method is gone: the strategies declare how a run is evaluated and the fitting unit carries it out, so the objective comes from there now. The tests themselves are unchanged -- they still check that a bad trial is pruned, that disabling the pruner completes every trial, and that a real model reports each epoch to its trial. Which partitions a trial records is still read off the strategy class, the same way the job reads it, so the declaration stays connected to what it produces. These four failures were not caught earlier because the verification runs had been narrowed to the directories this work was touching -- units, dag, spike, api and evaluation -- after the full suite was dropped for containing a test that builds the app against the real ~/.DashAI. Deselecting that one test was the right answer; shrinking the suite to what seemed relevant was not, and it is precisely the change that removes a caller elsewhere that this hides. Whole suite: 3682 passed, one test deselected. Co-Authored-By: Claude Opus 5 (1M context) --- .../test_optuna_pruning_integration.py | 30 ++++++++++++++----- .../back/optimizers/test_optuna_real_model.py | 30 ++++++++++++++----- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/tests/back/optimizers/test_optuna_pruning_integration.py b/tests/back/optimizers/test_optuna_pruning_integration.py index 0f08be0d7..9078cfb21 100644 --- a/tests/back/optimizers/test_optuna_pruning_integration.py +++ b/tests/back/optimizers/test_optuna_pruning_integration.py @@ -29,18 +29,32 @@ from DashAI.back.evaluation.holdout import HoldoutEvaluationStrategy from DashAI.back.models.base_model import BaseModel from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.units.fit_model_unit import FitModelUnit def _holdout_evaluate(model, input_dataset, output_dataset, metric): - """The real holdout evaluation path, on a strategy with no factory. - - `evaluate` reads which partitions its strategy scores, so it needs a real - instance rather than None for self. Building one through __init__ would - need a `ModelFactory` this test does not have, and does not need: the only - thing read off the instance is a class attribute. + """The real objective of a holdout search: one fit, then score validation. + + It used to be ``HoldoutEvaluationStrategy.evaluate``. The strategies now + only declare how a run is evaluated, and the unit that fits carries it out, + so the objective the optimizer measures comes from there. Which partitions + a trial records is still the strategy's declaration -- the job reads + SCORED_SPLITS off it and hands it over, minus the test partition, which a + trial may never score -- so it is read off the class here too. """ - strategy = HoldoutEvaluationStrategy.__new__(HoldoutEvaluationStrategy) - return strategy.evaluate(model, input_dataset, output_dataset, metric) + trial_splits = [ + split.name + for split in HoldoutEvaluationStrategy.SCORED_SPLITS + if split.name != "TEST" + ] + unit = FitModelUnit( + optimizer={"component": "", "params": {}}, + goal_metric="", + run_id=None, + artifact_prefix=None, + trial_splits=trial_splits, + ) + return unit._score_one_trial(model, input_dataset, output_dataset, metric) EPOCHS = 12 diff --git a/tests/back/optimizers/test_optuna_real_model.py b/tests/back/optimizers/test_optuna_real_model.py index 9eeb7dba7..0a49431d4 100644 --- a/tests/back/optimizers/test_optuna_real_model.py +++ b/tests/back/optimizers/test_optuna_real_model.py @@ -40,18 +40,32 @@ from DashAI.back.metrics.classification.accuracy import Accuracy from DashAI.back.models.mlp_image_classifier import MLPImageClassifier from DashAI.back.optimizers.optuna_optimizer import OptunaOptimizer +from DashAI.back.units.fit_model_unit import FitModelUnit def _holdout_evaluate(model, input_dataset, output_dataset, metric): - """The real holdout evaluation path, on a strategy with no factory. - - `evaluate` reads which partitions its strategy scores, so it needs a real - instance rather than None for self. Building one through __init__ would - need a `ModelFactory` this test does not have, and does not need: the only - thing read off the instance is a class attribute. + """The real objective of a holdout search: one fit, then score validation. + + It used to be ``HoldoutEvaluationStrategy.evaluate``. The strategies now + only declare how a run is evaluated, and the unit that fits carries it out, + so the objective the optimizer measures comes from there. Which partitions + a trial records is still the strategy's declaration -- the job reads + SCORED_SPLITS off it and hands it over, minus the test partition, which a + trial may never score -- so it is read off the class here too. """ - strategy = HoldoutEvaluationStrategy.__new__(HoldoutEvaluationStrategy) - return strategy.evaluate(model, input_dataset, output_dataset, metric) + trial_splits = [ + split.name + for split in HoldoutEvaluationStrategy.SCORED_SPLITS + if split.name != "TEST" + ] + unit = FitModelUnit( + optimizer={"component": "", "params": {}}, + goal_metric="", + run_id=None, + artifact_prefix=None, + trial_splits=trial_splits, + ) + return unit._score_one_trial(model, input_dataset, output_dataset, metric) EPOCHS = 3 From 8ce15ebee2e1bef8b7d5872cf35a65c31144003b Mon Sep 17 00:00:00 2001 From: Felipedino Date: Thu, 10 Sep 2026 22:23:38 -0300 Subject: [PATCH 28/28] Fix what the review found, and teach the audit about inheritance Five real findings, one of them hiding the others. **A search needs a tuner, not only a target.** Both evaluation strategies guarded on `self.optimizer and self.run_optimizable_parameters`; the units kept only the second half. `Run.optimizer_name` is a plain string and the wizard leaves it empty when no search is asked for, while the model may still declare a parameter optimizable -- a combination that has always meant "fit it once with the values given". It had become a lookup of the empty string in the registry, surfacing as "Metric is not compatible with the Task. ''", a message with nothing to do with what happened. Reproduced, fixed with a shared `_will_search`, and pinned by a test. **The nested unit was not being audited at all.** `_unit_class` matched only classes whose direct base is literally `BaseUnit`, so a unit that extends another unit fell out of every contract check -- and would have failed them, because its PROVIDES are written by the parent's body rather than its own. It is the same blindness a shared helper causes, arriving by inheritance instead: the audit reads one class's source. It now follows the lineage for declarations, context calls and config reads. 32 audited units became 33. **The inner splitter was resolved unconditionally**, so a run still carrying a nested configuration it no longer uses failed on a splitter it would never have touched. **The fold branch hardcoded `splits=["TEST"]`** where the holdout branch derives it from SCORED_SPLITS. Latent today -- no strategy excludes TEST -- but it is exactly the coupling this work exists to remove. And a docstring describing `{split: [scores]}` for something shaped `{split: {metric: [scores]}}`. Two findings were left alone, deliberately. `best_parameters` is published without being in PROVIDES, which is the already-declared limitation that there is no way to express an optional output; the new units repeat it rather than inventing an exception to it. And per-fold progress reporting is gone, which is a real regression: restoring it needs a callback in a unit's contract, and a runtime parameter the engine cannot supply makes the unit unusable as a node -- the static validator rejects it. Both are recorded rather than patched over. Whole suite: 3692 passed. Co-Authored-By: Claude Opus 5 (1M context) --- DashAI/back/job/model_job.py | 2 +- .../back/units/fit_model_over_folds_unit.py | 9 +- .../units/fit_model_over_nested_folds_unit.py | 12 +-- DashAI/back/units/fit_model_unit.py | 2 +- DashAI/back/units/fit_scope.py | 22 ++++- tests/back/units/test_fit_model_unit.py | 23 +++++ tests/back/units/test_unit_contracts.py | 84 +++++++++++++++++-- 7 files changed, 130 insertions(+), 24 deletions(-) diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index c105dee89..0f16ede98 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -294,7 +294,7 @@ def run( # uses. Whether there is anything to score is the # caller's to know: a session that reserved nothing # leaves that partition empty rather than absent. - if len(x[-1]["test"]) > 0: + if "TEST" in scored_splits and len(x[-1]["test"]) > 0: EvaluateModelUnit(run_id=run_id, splits=["TEST"])(ctx) except Exception as e: log.exception(e) diff --git a/DashAI/back/units/fit_model_over_folds_unit.py b/DashAI/back/units/fit_model_over_folds_unit.py index 662481002..b530df41b 100644 --- a/DashAI/back/units/fit_model_over_folds_unit.py +++ b/DashAI/back/units/fit_model_over_folds_unit.py @@ -96,7 +96,7 @@ def execute(self, ctx: ExecutionContext) -> None: plot_paths = [] try: - if optimizable_parameters: + if self._will_search(optimizable_parameters): # Every read of the context happens in this file rather than in # the shared helper: the contract audit parses it, so a require # moved out makes a declared key look unread. @@ -221,9 +221,10 @@ def _score_every_fold(self, model, x_folds, y_folds) -> dict: Returns ------- dict - ``{split name: [one score per fold]}``, in fold order, for whoever - aggregates them. A split with no metrics configured is absent - rather than present and empty: the two are different statements. + ``{split name: {metric name: [one score per fold]}}``, in fold + order, for whoever aggregates them. A split with no metrics + configured is absent rather than present and empty: the two are + different statements. """ scored_splits = self.config.get("scored_splits", TRIAL_SPLITS) fold_metrics: dict = {} diff --git a/DashAI/back/units/fit_model_over_nested_folds_unit.py b/DashAI/back/units/fit_model_over_nested_folds_unit.py index f55505da5..a97a99f83 100644 --- a/DashAI/back/units/fit_model_over_nested_folds_unit.py +++ b/DashAI/back/units/fit_model_over_nested_folds_unit.py @@ -102,11 +102,13 @@ def execute(self, ctx: ExecutionContext) -> None: y_folds = ctx.require("y_folds") optimizable_parameters = ctx.require("optimizable_parameters") - # Resolved outside the wrapper below so a splitter that cannot be built - # is reported as that, rather than as a training failure. - inner_splitter = self._resolve_inner_splitter() - - if optimizable_parameters: + if self._will_search(optimizable_parameters): + # Resolved here and not at the top: a run with nothing to search + # never carves an outer fold, so a session left carrying a nested + # configuration it no longer uses must not fail on it. And resolved + # outside the wrapper below, so a splitter that cannot be built is + # reported as that rather than as a training failure. + inner_splitter = self._resolve_inner_splitter() try: outer_fold_metrics = self._measure_every_outer_fold( model, diff --git a/DashAI/back/units/fit_model_unit.py b/DashAI/back/units/fit_model_unit.py index 018c2e0da..3e3d69b10 100644 --- a/DashAI/back/units/fit_model_unit.py +++ b/DashAI/back/units/fit_model_unit.py @@ -84,7 +84,7 @@ def execute(self, ctx: ExecutionContext) -> None: plot_paths = [] try: - if not optimizable_parameters: + if not self._will_search(optimizable_parameters): self._fit_kept_model(model, x, y) else: # Every read of the context happens here rather than in the diff --git a/DashAI/back/units/fit_scope.py b/DashAI/back/units/fit_scope.py index aa6fdba1d..e484cc26c 100644 --- a/DashAI/back/units/fit_scope.py +++ b/DashAI/back/units/fit_scope.py @@ -216,17 +216,31 @@ def _resolve_search(self): self._optimizer = optimizer return optimizer, goal_metric + def _will_search(self, optimizable_parameters) -> bool: + """Whether there is a search to run: something to tune, and a tuner. + + Both halves are needed. A run can name no optimizer at all -- the + column is a plain string and the wizard leaves it empty when the user + does not ask for a search -- while the model still declares a parameter + as optimizable, and that combination has always meant "fit it once with + the values given". Checking only the parameters turns it into a lookup + of the empty string in the registry, which fails with a message about + the metric being incompatible with the task. + """ + return bool(optimizable_parameters) and bool( + self.config["optimizer"]["component"] + ) + def _validate_search(self, optimizable_parameters) -> None: """Refuse an impossible search before anything observable happens. Handed the value rather than the context: the caller reads it with ``ctx.require`` and not ``ctx.get``, because an absent key means the model has not been built yet -- a call-order mistake, not "there is - nothing to optimize". Only an empty value, the key present and the - model declaring none, skips the checks below, so no registry lookup is - needed either. + nothing to optimize". A run with nothing to search skips the checks + below, so no registry lookup is needed either. """ - if not optimizable_parameters: + if not self._will_search(optimizable_parameters): return self._resolve_search() diff --git a/tests/back/units/test_fit_model_unit.py b/tests/back/units/test_fit_model_unit.py index e63d3c282..814835607 100644 --- a/tests/back/units/test_fit_model_unit.py +++ b/tests/back/units/test_fit_model_unit.py @@ -379,3 +379,26 @@ def test_the_search_is_handed_the_units_own_objective(tmp_path): finally: del di["component_registry"] del di["config"] + + +def test_a_run_with_no_optimizer_fits_once_even_if_a_parameter_is_optimizable(): + """Both halves are needed to call something a search: a tuner and a target. + + ``Run.optimizer_name`` is a plain string and the wizard leaves it empty when + the user does not ask for a search, while the model may still declare a + parameter as optimizable. That combination has always meant "fit it once + with the values given". Checking only the parameters turns it into a lookup + of the empty string in the registry, which surfaces as a complaint about + the metric being incompatible with the task -- a message with nothing to do + with what happened. + """ + model = _RecordingModel() + ctx = _fit_context(model, _HOLDOUT, _HOLDOUT) + ctx.put("optimizable_parameters", [("obj", "C", (0, 1), "number")]) + + unit = _unit(optimizer_name="", goal_metric="") + unit.validate(ctx) + unit(ctx) + + assert model.fits == [{"train": "x-train", "validation": "x-val"}] + assert not ctx.has("best_parameters") diff --git a/tests/back/units/test_unit_contracts.py b/tests/back/units/test_unit_contracts.py index 60bea26ac..439ce19a3 100644 --- a/tests/back/units/test_unit_contracts.py +++ b/tests/back/units/test_unit_contracts.py @@ -34,26 +34,85 @@ def _string_literals(node): } +def _all_classes(): + """Every class defined under ``units/``, by name. + + Built first so a unit that inherits from another unit can be resolved: the + audit reads source rather than importing, so a base class is just a name + until something maps it back to a definition. + """ + classes = {} + for path in _unit_modules(): + tree = ast.parse(path.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + classes[node.name] = node + return classes + + +_CLASSES = _all_classes() + + +def _is_unit(node): + """A class deriving from ``BaseUnit``, directly or through another unit.""" + for base in node.bases: + if not isinstance(base, ast.Name): + continue + if base.id == "BaseUnit": + return True + parent = _CLASSES.get(base.id) + if parent is not None and parent is not node and _is_unit(parent): + return True + return False + + def _unit_class(tree): for node in ast.walk(tree): - if isinstance(node, ast.ClassDef) and any( - isinstance(base, ast.Name) and base.id == "BaseUnit" for base in node.bases - ): + if isinstance(node, ast.ClassDef) and _is_unit(node): return node return None +def _lineage(cls): + """The class and the units it inherits from, nearest first. + + Everything below reads the whole lineage rather than one class body. A unit + that extends another one inherits both its declarations and the code that + honours them, so auditing only its own body would report that it promises + keys it never writes -- which is the same blindness a shared helper causes, + arriving by a different road. + """ + chain = [cls] + for base in cls.bases: + if not isinstance(base, ast.Name): + continue + parent = _CLASSES.get(base.id) + if parent is not None and parent is not cls and _is_unit(parent): + chain.extend(_lineage(parent)) + return chain + + def _declared(cls, name): - for node in cls.body: - if isinstance(node, ast.Assign) and any( - isinstance(t, ast.Name) and t.id == name for t in node.targets - ): - return _string_literals(node.value) + # Nearest declaration wins: a subclass that redeclares PROVIDES replaces + # what it inherited rather than adding to it, the way Python resolves it. + for ancestor in _lineage(cls): + for node in ancestor.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == name for t in node.targets + ): + return _string_literals(node.value) return set() def _context_calls(cls, methods): - """Every ``ctx.("key")`` literal inside the class.""" + """Every ``ctx.("key")`` literal in the class and what it extends.""" + keys = set() + for ancestor in _lineage(cls): + keys |= _context_calls_in(ancestor, methods) + return keys + + +def _context_calls_in(cls, methods): keys = set() for node in ast.walk(cls): if ( @@ -302,6 +361,13 @@ def _schema_fields(tree): def _config_reads(cls): + reads = set() + for ancestor in _lineage(cls): + reads |= _config_reads_in(ancestor) + return reads + + +def _config_reads_in(cls): """Every ``self.config["key"]`` and ``self.config.get("key")`` in the class.""" keys = set() for node in ast.walk(cls):