perf: lazy loading for TensorFlow trainers - CLI now instant

Author: codeMaestro78

mlcli/__init__.py

@@ -6,13 +6,23 @@
 """
 
 
-__version__="0.1.0"
-__author__="Devarshi Lalani"
-__licence__="MIT"
+__version__ = "0.1.0"
+__author__ = "Devarshi Lalani"
+__licence__ = "MIT"
 
 from mlcli.utils.registry import ModelRegistry
 
 # Global model registry instance
+registry = ModelRegistry()
 
-registry=ModelRegistry()
-__all__=["registry","__version__"]
+
+def _register_models():
+    """Register all models lazily without importing heavy dependencies."""
+    from mlcli.trainers import register_all_models
+    register_all_models()
+
+
+# Register models on first access
+_register_models()
+
+__all__ = ["registry", "__version__"]

mlcli/trainers/__init__.py

@@ -1,15 +1,108 @@
-"""Model trainers module."""
+"""Model trainers module with lazy loading to avoid slow imports."""
 
-from mlcli.trainers.base_trainer import  BaseTrainer
+from mlcli.trainers.base_trainer import BaseTrainer
+
+# Lazy imports - TensorFlow trainers are only loaded when accessed
+_LAZY_IMPORTS = {
+    "LogisticRegressionTrainer": "mlcli.trainers.logistic_trainer",
+    "SVMTrainer": "mlcli.trainers.svm_trainer",
+    "RFTrainer": "mlcli.trainers.rf_trainer",
+    "XGBTrainer": "mlcli.trainers.xgb_trainer",
+    "TFDNNTrainer": "mlcli.trainers.tf_dnn_trainer",
+    "TFCNNTrainer": "mlcli.trainers.tf_cnn_trainer",
+    "TFRNNTrainer": "mlcli.trainers.tf_rnn_trainer",
+}
+
+# Pre-register models without importing heavy dependencies
+_MODEL_METADATA = {
+    "logistic_regression": {
+        "class": "LogisticRegressionTrainer",
+        "module": "mlcli.trainers.logistic_trainer",
+        "description": "Logistic Regression classifier with L2 regularization",
+        "framework": "sklearn",
+        "model_type": "classification",
+    },
+    "svm": {
+        "class": "SVMTrainer",
+        "module": "mlcli.trainers.svm_trainer",
+        "description": "Support Vector Machine with RBF/Linear/Poly kernels",
+        "framework": "sklearn",
+        "model_type": "classification",
+    },
+    "random_forest": {
+        "class": "RFTrainer",
+        "module": "mlcli.trainers.rf_trainer",
+        "description": "Random Forest ensemble classifier",
+        "framework": "sklearn",
+        "model_type": "classification",
+    },
+    "xgboost": {
+        "class": "XGBTrainer",
+        "module": "mlcli.trainers.xgb_trainer",
+        "description": "XGBoost gradient boosting classifier",
+        "framework": "xgboost",
+        "model_type": "classification",
+    },
+    "tf_dnn": {
+        "class": "TFDNNTrainer",
+        "module": "mlcli.trainers.tf_dnn_trainer",
+        "description": "Tensorflow Dense Feedforward Neural Network",
+        "framework": "tensorflow",
+        "model_type": "classification",
+    },
+    "tf_cnn": {
+        "class": "TFCNNTrainer",
+        "module": "mlcli.trainers.tf_cnn_trainer",
+        "description": "TensorFlow Convolutional Neural Network for image classification",
+        "framework": "tensorflow",
+        "model_type": "classification",
+    },
+    "tf_rnn": {
+        "class": "TFRNNTrainer",
+        "module": "mlcli.trainers.tf_rnn_trainer",
+        "description": "TensorFlow RNN/LSTM/GRU for sequence classification",
+        "framework": "tensorflow",
+        "model_type": "classification",
+    },
+}
+
+
+def __getattr__(name: str):
+    """Lazy import trainers only when accessed."""
+    if name in _LAZY_IMPORTS:
+        import importlib
+        module = importlib.import_module(_LAZY_IMPORTS[name])
+        return getattr(module, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def register_all_models():
+    """Register all models in the registry without importing heavy modules."""
+    from mlcli import registry
+    
+    for model_name, meta in _MODEL_METADATA.items():
+        if not registry.is_registered(model_name):
+            # Register with lazy loader
+            registry.register_lazy(
+                name=model_name,
+                module_path=meta["module"],
+                class_name=meta["class"],
+                description=meta["description"],
+                framework=meta["framework"],
+                model_type=meta["model_type"],
+            )
+
+
+def get_trainer_class(model_type: str):
+    """Get trainer class by model type, importing only when needed."""
+    if model_type not in _MODEL_METADATA:
+        raise ValueError(f"Unknown model type: {model_type}")
+    
+    import importlib
+    meta = _MODEL_METADATA[model_type]
+    module = importlib.import_module(meta["module"])
+    return getattr(module, meta["class"])
 
-# Import all trainers to trigger auto-registration
-from mlcli.trainers.logistic_trainer import LogisticRegressionTrainer
-from mlcli.trainers.svm_trainer import SVMTrainer
-from mlcli.trainers.rf_trainer import RFTrainer
-from mlcli.trainers.xgb_trainer import XGBTrainer
-from mlcli.trainers.tf_dnn_trainer import TFDNNTrainer
-from mlcli.trainers.tf_cnn_trainer import TFCNNTrainer
-from mlcli.trainers.tf_rnn_trainer import TFRNNTrainer
 
 __all__ = [
     "BaseTrainer",
@@ -20,4 +113,6 @@
     "TFDNNTrainer",
     "TFCNNTrainer",
     "TFRNNTrainer",
+    "register_all_models",
+    "get_trainer_class",
 ]

mlcli/utils/registry.py

@@ -3,6 +3,7 @@
 
 Provides decorator-based auto-registration for all trainer classes,
 enabling dynamic model discovery and instantiation from configuration.
+Supports lazy loading to avoid importing heavy dependencies like TensorFlow.
 """
 
 from typing import Dict, Type, Optional, List, Any
@@ -16,15 +17,18 @@ class ModelRegistry:
     Central registry for all model trainers.
 
     Maps model type strings (e.g., 'logistic_regression') to their corresponding
-    trainer class implementations. Supports automatic registration via decorator.
+    trainer class implementations. Supports automatic registration via decorator
+    and lazy loading for heavy dependencies.
     """
 
     def __init__(self) -> None:
         """Initialize empty registry."""
-        self._registry :Dict[str,Type]={}
-        self._metadata : Dict[str,Dict[str,Any]]={}
+        self._registry: Dict[str, Type] = {}
+        self._lazy_registry: Dict[str, Dict[str, str]] = {}  # For lazy loading
+        self._metadata: Dict[str, Dict[str, Any]] = {}
 
-    def register(self,name:str,trainer_class:Type,description:str="",framework:str="unknown",model_type:str="unknown")->None:
+    def register(self, name: str, trainer_class: Type, description: str = "",
+                 framework: str = "unknown", model_type: str = "unknown") -> None:
         """
         Register a trainer class with metadata.
 
@@ -42,17 +46,64 @@ def register(self,name:str,trainer_class:Type,description:str="",framework:str="
         if name in self._registry:
             logger.warning(f"Model '{name}' is already registered. Overwriting")
 
-        self._registry[name]= trainer_class
-        self._metadata[name]= {
-            "description":description,
-            "framework":framework,
-            "model_type":model_type,
-            "class_name":trainer_class.__name__
+        self._registry[name] = trainer_class
+        self._metadata[name] = {
+            "description": description,
+            "framework": framework,
+            "model_type": model_type,
+            "class_name": trainer_class.__name__
         }
 
-        logger.debug(f"Registered model:{name}->{trainer_class.__name__}")
+        logger.debug(f"Registered model: {name} -> {trainer_class.__name__}")
 
-    def get(self,name:str)->Optional[Type]:
+    def register_lazy(self, name: str, module_path: str, class_name: str,
+                      description: str = "", framework: str = "unknown",
+                      model_type: str = "unknown") -> None:
+        """
+        Register a trainer for lazy loading (doesn't import the module yet).
+
+        Args:
+            name: Unique identifier for the model
+            module_path: Full module path (e.g., 'mlcli.trainers.tf_dnn_trainer')
+            class_name: Class name to import from module
+            description: Human-readable description
+            framework: ML framework name
+            model_type: Type of model
+        """
+        self._lazy_registry[name] = {
+            "module_path": module_path,
+            "class_name": class_name,
+        }
+        self._metadata[name] = {
+            "description": description,
+            "framework": framework,
+            "model_type": model_type,
+            "class_name": class_name,
+            "lazy": True,
+        }
+        logger.debug(f"Registered lazy model: {name} -> {module_path}.{class_name}")
+
+    def _resolve_lazy(self, name: str) -> Optional[Type]:
+        """Resolve a lazy-registered model by importing it."""
+        if name not in self._lazy_registry:
+            return None
+        
+        import importlib
+        lazy_info = self._lazy_registry[name]
+        module = importlib.import_module(lazy_info["module_path"])
+        trainer_class = getattr(module, lazy_info["class_name"])
+        
+        # Move from lazy to regular registry
+        self._registry[name] = trainer_class
+        del self._lazy_registry[name]
+        
+        # Update metadata
+        if name in self._metadata:
+            self._metadata[name]["lazy"] = False
+        
+        return trainer_class
+
+    def get(self, name: str) -> Optional[Type]:
         """
         Retrieve a trainer class by name.
 
@@ -62,9 +113,17 @@ def get(self,name:str)->Optional[Type]:
         Returns:
             Trainer class or None if not found
         """
-        return self._registry.get(name)
-
-    def get_trainer(self,name:str,**kwargs)->Any:
+        # Check regular registry first
+        if name in self._registry:
+            return self._registry.get(name)
+        
+        # Try lazy loading
+        if name in self._lazy_registry:
+            return self._resolve_lazy(name)
+        
+        return None
+
+    def get_trainer(self, name: str, **kwargs) -> Any:
         """
         Instantiate a trainer by name.
 
@@ -78,21 +137,23 @@ def get_trainer(self,name:str,**kwargs)->Any:
         Raises:
             KeyError: If model name not found in registry
         """
-        trainer_class=self.get(name)
+        trainer_class = self.get(name)
         if trainer_class is None:
-            available=", ".join(self.list_models())
-            raise KeyError(f"Model '{name}' not found in registry." f"Available models: {available}")
+            available = ", ".join(self.list_models())
+            raise KeyError(f"Model '{name}' not found in registry. "
+                          f"Available models: {available}")
 
         return trainer_class(**kwargs)
 
-    def list_models(self)->List[str]:
+    def list_models(self) -> List[str]:
         """
-        Get list of all registered model names.
+        Get list of all registered model names (including lazy).
 
         Returns:
             List of model identifiers
         """
-        return sorted(self._registry.keys())
+        all_models = set(self._registry.keys()) | set(self._lazy_registry.keys())
+        return sorted(all_models)
 
 
     def get_metadata(self,name:str)->Optional[Dict[str,Any]]:
@@ -127,17 +188,17 @@ def get_models_by_framework(self,framework:str)->List[str]:
         """
         return [name for name,meta in self._metadata.items() if meta.get("framework")==framework]
 
-    def is_registered(self,name:str)->bool:
+    def is_registered(self, name: str) -> bool:
         """
-        Check if a model is registered.
+        Check if a model is registered (including lazy).
 
         Args:
             name: Model identifier
 
         Returns:
             True if registered, False otherwise
         """
-        return name in self._registry
+        return name in self._registry or name in self._lazy_registry
 
     def unregister(self,name:str)->bool:
         """