feat: conditional framework extras (#176)

* feat: make framework dependencies optional

* fix: treat keras as core framework

* fix: treat keras as core in retrain

Author: marcellodebernardi

Dockerfile

@@ -33,13 +33,20 @@ RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir poetry && \
     poetry config virtualenvs.create false
 
-# Install CPU-only PyTorch (no CUDA/NVIDIA packages - quickstart image includes torch)
-# Install before dependency copy so pyproject changes don't invalidate this layer.
-RUN pip install --no-cache-dir torch==2.7.1 --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
-
-# Install main dependencies + AWS support (no Spark provider yet)
+# Install large stable dependencies before poetry to maximize build cache reuse.
+# INSTALL_PYTORCH controls whether CPU-only PyTorch is installed.
+ARG INSTALL_PYTORCH="true"
+RUN if [ "$INSTALL_PYTORCH" = "true" ]; then \
+        pip install --no-cache-dir torch==2.7.1 \
+            --index-url https://download.pytorch.org/whl/cpu \
+            --extra-index-url https://pypi.org/simple; \
+    fi
+
+# Install main dependencies + optional framework extras (no Spark provider yet)
+# XGBoost is core; extras cover optional frameworks.
+ARG POETRY_EXTRAS="aws catboost"
 COPY pyproject.toml poetry.lock /code/
-RUN poetry install --only=main --no-root --extras aws
+RUN poetry install --only=main --no-root --extras "${POETRY_EXTRAS}"
 
 # Application code
 COPY plexe/ /code/plexe/

Makefile

@@ -250,7 +250,6 @@ test-full: build
 			--enable-final-evaluation
 	@echo "✅ Full test passed!"
 
-
 # ============================================
 # Example Datasets
 # ============================================
@@ -321,6 +320,7 @@ build:
 		-f Dockerfile .
 	@echo "✅ Build complete: plexe:py$(PYTHON_VERSION)"
 
+
 # Build Databricks variant
 .PHONY: build-databricks
 build-databricks:

README.md

@@ -178,9 +178,18 @@ See [`plexe/integrations/base.py`](plexe/integrations/base.py) for the full inte
 
 ### 3.1. Installation Options
 ```bash
-pip install plexe                    # Core (XGBoost, CatBoost, LightGBM, Keras, PyTorch, scikit-learn)
-pip install plexe[pyspark]           # + Local PySpark execution
-pip install plexe[aws]               # + S3 storage support (boto3)
+pip install plexe                    # Core (XGBoost, Keras, scikit-learn)
+```
+
+You can add optional dependencies either by framework or by task grouping:
+- Framework extras: `catboost`, `lightgbm`, `pytorch`
+- Task extras: `tabular` (CatBoost + LightGBM), `vision` (PyTorch)
+- Platform extras: `pyspark`, `aws`
+
+Examples:
+```bash
+pip install "plexe[tabular,pyspark]"   # tabular stack + local PySpark
+pip install "plexe[pytorch,aws]"       # explicit framework + S3 support
 ```
 
 Requires Python >= 3.10, < 3.13.

plexe/config.py

@@ -50,19 +50,28 @@ class ModelType:
     PYTORCH = "pytorch"
 
 
-# Default model types (enabled by default, user can override via --allowed-model-types)
-def _is_module_available(module_name: str) -> bool:
-    return importlib.util.find_spec(module_name) is not None
+def detect_installed_frameworks() -> list[str]:
+    """Detect which ML frameworks are installed and importable.
 
+    XGBoost and Keras are always available (core dependencies). Other frameworks are
+    checked via importlib.util.find_spec() which is fast and has no side effects.
+    """
+    available = [ModelType.XGBOOST, ModelType.KERAS]  # Always available (core dependencies)
 
-DEFAULT_MODEL_TYPES = [
-    ModelType.XGBOOST,
-    ModelType.CATBOOST,
-    ModelType.LIGHTGBM,
-    ModelType.KERAS,
-]
-if _is_module_available("torch"):
-    DEFAULT_MODEL_TYPES.append(ModelType.PYTORCH)
+    _OPTIONAL_FRAMEWORKS = [
+        ("catboost", ModelType.CATBOOST),
+        ("lightgbm", ModelType.LIGHTGBM),
+        ("torch", ModelType.PYTORCH),
+    ]
+    for module_name, model_type in _OPTIONAL_FRAMEWORKS:
+        if importlib.util.find_spec(module_name) is not None:
+            available.append(model_type)
+
+    return available
+
+
+# Default model types (auto-detected from installed packages, user can override via --allowed-model-types)
+DEFAULT_MODEL_TYPES = detect_installed_frameworks()
 
 # Task-compatible model types based on data layout
 # Maps DataLayout enum to compatible model types

plexe/helpers.py

@@ -17,7 +17,13 @@
 if TYPE_CHECKING:
     from pyspark.sql import SparkSession
 
-from plexe.config import ModelType, StandardMetric, DEFAULT_MODEL_TYPES, TASK_COMPATIBLE_MODELS
+from plexe.config import (
+    DEFAULT_MODEL_TYPES,
+    TASK_COMPATIBLE_MODELS,
+    ModelType,
+    StandardMetric,
+    detect_installed_frameworks,
+)
 from plexe.models import DataLayout
 
 logger = logging.getLogger(__name__)
@@ -47,6 +53,13 @@ def select_viable_model_types(data_layout: DataLayout, selected_frameworks: list
         selected = DEFAULT_MODEL_TYPES
         logger.info(f"Using default model types: {selected}")
     else:
+        # Validate that user-selected frameworks are actually installed
+        installed = detect_installed_frameworks()
+        not_installed = [f for f in selected_frameworks if f not in installed]
+        if not_installed:
+            raise ValueError(
+                f"Requested model types {not_installed} are not installed. " f"Installed frameworks: {installed}"
+            )
         selected = selected_frameworks
         logger.info(f"User-selected model types: {selected}")
 

plexe/retrain.py

@@ -5,15 +5,19 @@
 """
 
 import inspect
+import os
 import json
 import logging
 import shutil
 import tarfile
 from pathlib import Path
 
 import joblib
-import keras
 import pandas as pd
+
+# Ensure keras uses TensorFlow backend even when retrain is invoked directly.
+os.environ.setdefault("KERAS_BACKEND", "tensorflow")
+import keras
 import yaml
 from sklearn.model_selection import train_test_split
 from xgboost import XGBClassifier, XGBRegressor
@@ -209,7 +213,10 @@ def retrain_model(
             logger.info("Created new untrained XGBoost model from architecture")
 
         elif model_type == "catboost":
-            from catboost import CatBoostClassifier, CatBoostRegressor
+            try:
+                from catboost import CatBoostClassifier, CatBoostRegressor
+            except ImportError:
+                raise RetrainingError("CatBoost is required for retraining CatBoost models but is not installed")
 
             # CatBoost uses native format — try classifier first, then regressor
             try:
@@ -229,7 +236,10 @@ def retrain_model(
             logger.info("Created new untrained CatBoost model from architecture")
 
         elif model_type == "lightgbm":
-            from lightgbm import LGBMClassifier, LGBMRanker
+            try:
+                from lightgbm import LGBMClassifier, LGBMRanker
+            except ImportError:
+                raise RetrainingError("LightGBM is required for retraining LightGBM models but is not installed")
 
             trained_model = joblib.load(original_model_path)
             params = trained_model.get_params()

plexe/tools/submission.py

@@ -323,7 +323,10 @@ def save_model(model: Any) -> str:
                 logger.info(f"XGBoost model validated: {type(model).__name__}")
 
             elif model_type == "catboost":
-                from catboost import CatBoostClassifier, CatBoostRegressor
+                try:
+                    from catboost import CatBoostClassifier, CatBoostRegressor
+                except ImportError:
+                    raise ValueError("CatBoost is required for CatBoost models but is not installed")
 
                 if not isinstance(model, CatBoostClassifier | CatBoostRegressor):
                     error_msg = f"Expected CatBoostClassifier or CatBoostRegressor, got {type(model)}"
@@ -333,7 +336,10 @@ def save_model(model: Any) -> str:
                 logger.info(f"CatBoost model validated: {type(model).__name__}")
 
             elif model_type == "lightgbm":
-                from lightgbm import LGBMClassifier, LGBMRegressor, LGBMRanker
+                try:
+                    from lightgbm import LGBMClassifier, LGBMRegressor, LGBMRanker
+                except ImportError:
+                    raise ValueError("LightGBM is required for LightGBM models but is not installed")
 
                 if not isinstance(model, LGBMClassifier | LGBMRegressor | LGBMRanker):
                     error_msg = f"Expected LGBMClassifier, LGBMRegressor, or LGBMRanker, got {type(model)}"

plexe/workflow.py

@@ -482,6 +482,15 @@ def build_model(
             logger.info("Skipping Phase 6 (already completed)")
             final_package_dir = work_dir / "model"
 
+        summary_journal = journal or context.scratch.get("_search_journal")
+        if summary_journal and summary_journal.nodes:
+            total_attempts = len(summary_journal.nodes)
+            logger.info(
+                "Search summary: %s/%s solutions succeeded",
+                summary_journal.successful_attempts,
+                total_attempts,
+            )
+
         logger.info(f"Model building complete! Validation performance: {final_metrics}")
         logger.info(f"Final model package: {final_package_dir}")