Skip to content

Refactor datasetjob with units - #804

Open
Felipedino wants to merge 1 commit into
feat/atom-expl-predict-explorfrom
feat/atom_datasetjob
Open

Refactor datasetjob with units#804
Felipedino wants to merge 1 commit into
feat/atom-expl-predict-explorfrom
feat/atom_datasetjob

Conversation

@Felipedino

Copy link
Copy Markdown
Collaborator

This pull request introduces several improvements and refactorings across the backend, focusing on explainability model integration, job unit modularization, and bug fixes for dataset status handling. The most significant changes include the addition of a robust wrapper for SHAP model prediction to improve compatibility, a major refactor of the converter job to use new modular "unit" classes, and the registration of these units in the system's component registry. There are also minor dependency and bug fixes.

Explainability Integration and SHAP Compatibility:

  • Added a new as_shap_predictor function in model_input.py to wrap model prediction methods as plain functions, improving compatibility with SHAP and preventing errors with read-only properties in certain models (e.g., LightGBM, XGBoost). All SHAP explainer classes now use this wrapper instead of passing bound methods directly. [1] [2] [3] [4] [5] [6] [7]

Job and Unit System Refactor:

  • Refactored converter_job.py to remove the internal dataset transformation logic and instead leverage new modular "unit" classes (ApplyConverterUnit, LoadDatasetUnit, SaveDatasetUnit, etc.), streamlining the job's structure and reducing code duplication. [1] [2]
  • Registered all new unit classes in initial_components.py so they are available system-wide for orchestration and dependency injection. [1] [2]

Bug Fixes and Dependency Updates:

  • Fixed a bug in models.py where dataset status methods attempted to set start_time and end_time attributes that do not exist for the Dataset table, preventing unnecessary AttributeErrors.
  • Cleaned up dependencies in the preview_manual_prediction API endpoint by removing the unused component_registry parameter. [1] [2]

- 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.
Copilot AI lite review requested due to automatic review settings August 10, 2026 14:11

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR refactors backend job execution toward a composable “Unit” abstraction, improves SHAP compatibility by avoiding passing bound model methods, and tightens dataset/explorer status handling. It also adds a substantial regression/contract test suite around units and moves test dataset caches out of the repository tree.

Changes:

  • Introduce an ExecutionContext + BaseUnit contract system, register units in the ComponentRegistry, and refactor ModelJob, ConverterJob, and ExplorerJob to orchestrate via units.
  • Add as_shap_predictor() and update SHAP explainers to pass a plain callable; add regression tests covering read-only feature_names_in_ wrappers.
  • Improve backend/test hygiene: dataset status bugfix, remove unused endpoint dependency, and route HF dataset cache writes to a shared scratch directory cleared per test session.

Reviewed changes

Copilot reviewed 84 out of 85 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
tests/back/units/test_unit_contracts.py Adds an AST-based audit test to enforce unit context REQUIRES/PROVIDES contract correctness.
tests/back/units/test_save_dataset_unit.py Adds contract tests for SaveDatasetUnit.
tests/back/units/test_load_dataset_unit.py Adds contract tests for LoadDatasetUnit, including notebook/dataset branches and error cases.
tests/back/units/test_fit_model_unit.py Adds regression tests for FitModelUnit.validate() contract behavior and instance-local caching.
tests/back/units/test_evaluate_model_unit.py Adds regression test ensuring run_id is required (no silent no-op).
tests/back/units/test_converter_fit_transform_split.py Adds tests for split fit/transform converter workflow and contract expectations.
tests/back/units/test_context.py Adds thorough tests for ExecutionContext ref/cached semantics and copy guarantees.
tests/back/units/test_build_model_unit.py Adds tests for per-instance model class memoization vs context-global caching.
tests/back/units/test_base_unit.py Adds tests for BaseUnit contract enforcement, validate-before-execute behavior, and registry typing.
tests/back/units/init.py Establishes units test package marker.
tests/back/tasks/test_tasks.py Switches dataloader temp paths to shared scratch directory.
tests/back/scratch.py Introduces per-PID scratch dir utilities for tests and best-effort cleanup.
tests/back/models/test_tabular_class_models.py Moves HF cache writes to scratch dir.
tests/back/models/test_modernbert_transformer.py Moves HF cache writes to scratch dir.
tests/back/models/test_distilbert_transformer.py Moves HF cache writes to scratch dir.
tests/back/models/test_deberta_v3_transformer.py Moves HF cache writes to scratch dir.
tests/back/models/test_bow_text_class_model.py Moves HF cache writes to scratch dir.
tests/back/explainers/test_task_explainers.py Moves HF cache writes to scratch dir.
tests/back/explainers/test_shap_predictor_handover.py Adds regression tests for SHAP method handover and read-only feature_names_in_.
tests/back/explainers/test_new_explainers.py Moves HF cache writes to scratch dir.
tests/back/explainers/test_lib_explainers.py Moves HF cache writes to scratch dir.
tests/back/explainers/test_explainers.py Moves HF cache writes to scratch dir.
tests/back/dataloaders/test_dashai_dataset.py Moves HF cache writes to scratch dir.
tests/back/dataloaders/base_tabular_dataloader_tests.py Moves HF cache writes to scratch dir.
tests/back/conftest.py Adds session-level scratch cleanup fixture.
tests/back/api/test_units_api.py Adds API coverage asserting units are registered and schema-exposed as components.
tests/back/api/test_explorer_job.py Adds explicit end-to-end regression suite for ExplorerJob behavior and status transitions.
DashAI/back/units/transform_dataset_unit.py Adds TransformDatasetUnit for applying an already-fitted converter to a dataset.
DashAI/back/units/save_prediction_unit.py Adds SavePredictionUnit to persist predictions and publish results path.
DashAI/back/units/save_model_unit.py Adds SaveModelUnit to persist trained models and publish model path.
DashAI/back/units/save_exploration_unit.py Adds SaveExplorationUnit to persist exploration artifacts and publish exploration path.
DashAI/back/units/save_dataset_unit.py Adds SaveDatasetUnit to persist datasets back to their loaded path.
DashAI/back/units/save_dataset_to_path_unit.py Adds SaveDatasetToPathUnit for saving datasets to a configured destination.
DashAI/back/units/run_exploration_unit.py Adds RunExplorationUnit to instantiate and run an explorer and publish result + explorer instance.
DashAI/back/units/prepare_explanation_data_unit.py Adds PrepareExplanationDataUnit to replay recorded split indexes for explainability flows.
DashAI/back/units/prepare_and_split_unit.py Adds PrepareAndSplitUnit to prepare/split datasets for training and publish inputs/targets/splits.
DashAI/back/units/predict_unit.py Adds PredictUnit to run model prediction and decode outputs via the task.
DashAI/back/units/load_uploaded_dataset_unit.py Adds LoadUploadedDatasetUnit to parse a dataset from an uploaded file/URL.
DashAI/back/units/load_training_dataset_unit.py Adds LoadTrainingDatasetUnit to load training dataset + types for prediction decoding/saving.
DashAI/back/units/load_trained_model_unit.py Adds LoadTrainedModelUnit to restore models by run id using class-level load.
DashAI/back/units/load_run_model_unit.py Adds LoadRunModelUnit to preserve the legacy “instantiate then load” behavior for explainers.
DashAI/back/units/load_dataset_unit.py Adds LoadDatasetUnit to load stored datasets/notebook working copies and publish id/path.
DashAI/back/units/load_datafile_dataset_unit.py Adds LoadDatafileDatasetUnit to read datasets from completed hub downloads.
DashAI/back/units/infer_dataset_types_unit.py Adds InferDatasetTypesUnit to publish inferred or existing dataset types.
DashAI/back/units/generate_global_explanation_unit.py Adds GenerateGlobalExplanationUnit to run global explainers and store artifacts.
DashAI/back/units/fit_model_unit.py Adds FitModelUnit to train models and optionally run HPO, publishing plots and best params.
DashAI/back/units/fit_converter_unit.py Adds FitConverterUnit to fit converters without transforming data, publishing fitted converter.
DashAI/back/units/explanation_artifacts.py Adds helper functions for explainer instantiation and explanation artifact dumping.
DashAI/back/units/evaluate_model_unit.py Adds EvaluateModelUnit for idempotent final-metric computation per split.
DashAI/back/units/context.py Introduces ExecutionContext with ref/cache halves and copy-on-read guarantees for refs.
DashAI/back/units/compute_dataset_metadata_unit.py Adds unit to compute/strip dataset metadata according to configuration.
DashAI/back/units/build_model_unit.py Adds BuildModelUnit to resolve model class, validate downloads, and build ModelFactory outputs.
DashAI/back/units/build_manual_input_unit.py Adds BuildManualInputUnit to construct a prediction dataset from user-entered rows.
DashAI/back/units/build_local_explainer_unit.py Adds BuildLocalExplainerUnit for local explainer instantiation.
DashAI/back/units/build_global_explainer_unit.py Adds BuildGlobalExplainerUnit for global explainer instantiation.
DashAI/back/units/base_unit.py Introduces BaseUnit contract enforcement, validate-before-execute, and registry TYPE.
DashAI/back/units/apply_dataset_schema_unit.py Adds unit to rename/cast dataset columns based on inferred types + optional renames.
DashAI/back/units/apply_converter_unit.py Adds fused “fit+transform” converter unit, publishing dataset + fitted converter.
DashAI/back/units/init.py Exposes unit primitives (BaseUnit, ExecutionContext, UnitContractError).
DashAI/back/job/model_job.py Refactors ModelJob to orchestrate data prep/build/fit/eval/save via units and shared context.
DashAI/back/job/explorer_job.py Refactors ExplorerJob to load dataset/run exploration/save via units with improved error status handling.
DashAI/back/job/converter_job.py Refactors ConverterJob to load/apply/save via units with contract-driven validation.
DashAI/back/initial_components.py Registers all new units as components so they’re discoverable via the registry/API.
DashAI/back/explainability/model_input.py Adds as_shap_predictor() wrapper to avoid passing bound methods into SHAP.
DashAI/back/explainability/explainers/regression_kernel_shap.py Updates explainer to use as_shap_predictor(self.model) for SHAP model handover.
DashAI/back/explainability/explainers/kernel_shap.py Updates explainer to use as_shap_predictor(self.model) for SHAP model handover.
DashAI/back/explainability/explainers/contrastive_shap.py Updates explainer to use as_shap_predictor(self.model) for SHAP model handover.
DashAI/back/dependencies/database/models.py Fixes dataset status setters to stop assigning non-existent start/end time columns.
DashAI/back/api/api_v1/endpoints/predict.py Removes unused component_registry dependency from preview_manual_prediction.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +57 to +60
raise JobError(
f"""Unable to find the {scope} explainer with name
{explainer_name} in registry.""",
) from e
Comment on lines +136 to +140
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"]()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants