⚡ perf(cli): faster CLI startup via lazy imports

src/anomalib/cli/cli.py

@@ -15,31 +15,25 @@
 from typing import Any
 
 from jsonargparse import ActionConfigFile, ArgumentParser, Namespace
-from jsonargparse._actions import _ActionSubCommands
 from rich import traceback
 
+try:
+    from jsonargparse._actions import _ActionSubCommands
+except ImportError:
+    from jsonargparse._subcommands import ActionSubCommands as _ActionSubCommands
+
 from anomalib import __version__
-from anomalib.cli.pipelines import PIPELINE_REGISTRY, pipeline_subcommands, run_pipeline
-from anomalib.cli.utils.help_formatter import CustomHelpFormatter, get_short_docstring
-from anomalib.cli.utils.openvino import add_openvino_export_arguments
-from anomalib.loggers import configure_logger
+from anomalib.cli.utils.help_formatter import CustomHelpFormatter
 
 traceback.install()
 logger = logging.getLogger("anomalib.cli")
 
-_LIGHTNING_AVAILABLE = True
-try:
-    from lightning.pytorch import Trainer
-    from torch.utils.data import DataLoader, Dataset
 
-    from anomalib.data import AnomalibDataModule
-    from anomalib.engine import Engine
-    from anomalib.models import AnomalibModule
-    from anomalib.utils.config import update_config
+def _check_lightning_available() -> bool:
+    """Check if Lightning and its dependencies are installed (without importing them)."""
+    import importlib.util
 
-except ImportError as error:
-    _LIGHTNING_AVAILABLE = False
-    logger.warning(f"Import failed: {error}. Please install the required dependencies.")
+    return importlib.util.find_spec("lightning") is not None and importlib.util.find_spec("torch") is not None
 
 
 class AnomalibCLI:
@@ -69,19 +63,48 @@ class AnomalibCLI:
         provided via both files and command line arguments simultaneously.
     """
 
+    _TRAINER_SUBCOMMAND_DESCRIPTIONS: dict[str, str] = {
+        "fit": "Runs the full optimization routine.",
+        "validate": "Perform one evaluation epoch over the validation set.",
+        "test": "Perform one evaluation epoch over the test set.",
+    }
+
     def __init__(self, args: Sequence[str] | None = None, run: bool = True) -> None:
+        self._lightning_available = _check_lightning_available()
+        self._selected_subcommand = self._sniff_subcommand(args)
         self.parser = self.init_parser()
         self.subcommand_parsers: dict[str, ArgumentParser] = {}
         self.subcommand_method_arguments: dict[str, list[str]] = {}
         self.add_subcommands()
         self.config = self.parser.parse_args(args=args)
         self.subcommand = self.config["subcommand"]
-        if _LIGHTNING_AVAILABLE:
+        if self._lightning_available:
             self.before_instantiate_classes()
             self.instantiate_classes()
         if run:
             self._run_subcommand()
 
+    # Flags whose next token is a value, not a subcommand.
+    _OPTIONS_WITH_VALUE = frozenset({"-c", "--config"})
+
+    @staticmethod
+    def _sniff_subcommand(args: Sequence[str] | None) -> str | None:
+        """Peek at args to identify the subcommand without full parsing."""
+        import sys
+
+        tokens = list(args) if args is not None else sys.argv[1:]
+        skip_next = False
+        for token in tokens:
+            if skip_next:
+                skip_next = False
+                continue
+            if token in AnomalibCLI._OPTIONS_WITH_VALUE:
+                skip_next = True
+                continue
+            if not token.startswith("-"):
+                return token
+        return None
+
     @staticmethod
     def init_parser(**kwargs) -> ArgumentParser:
         """Method that instantiates the argument parser."""
@@ -120,41 +143,47 @@ def add_subcommands(self, **kwargs) -> None:
         # Extra subcommand: install
         self._set_install_subcommand(parser_subcommands)
 
-        if not _LIGHTNING_AVAILABLE:
-            # If environment is not configured to use pl, do not add a subcommand for Engine.
+        if not self._lightning_available:
             return
 
-        # Add Trainer subcommands
+        selected = self._selected_subcommand
+
+        # Add Trainer subcommands (fit, validate, test)
         for subcommand in self.subcommands():
             sub_parser = self.init_parser(**kwargs)
-
-            fn = getattr(Trainer, subcommand)
-            # extract the first line description in the docstring for the subcommand help message
-            description = get_short_docstring(fn)
-            subparser_kwargs = kwargs.get(subcommand, {})
-            subparser_kwargs.setdefault("description", description)
-
+            description = self._TRAINER_SUBCOMMAND_DESCRIPTIONS.get(subcommand, "")
             self.subcommand_parsers[subcommand] = sub_parser
             parser_subcommands.add_subcommand(subcommand, sub_parser, help=description)
-            self.add_trainer_arguments(sub_parser, subcommand)
+            if selected == subcommand:
+                self.add_trainer_arguments(sub_parser, subcommand)
 
-        # Add anomalib subcommands
+        # Add anomalib subcommands (train, predict, export)
         for subcommand in self.anomalib_subcommands():
             sub_parser = self.init_parser(**kwargs)
-
             self.subcommand_parsers[subcommand] = sub_parser
             parser_subcommands.add_subcommand(
                 subcommand,
                 sub_parser,
                 help=self.anomalib_subcommands()[subcommand]["description"],
             )
-            # add arguments to subcommand
-            getattr(self, f"add_{subcommand}_arguments")(sub_parser)
+            if selected == subcommand:
+                getattr(self, f"add_{subcommand}_arguments")(sub_parser)
 
         # Add pipeline subcommands
-        if PIPELINE_REGISTRY is not None:
-            for subcommand, value in pipeline_subcommands().items():
-                sub_parser = PIPELINE_REGISTRY[subcommand].get_parser()
+        from anomalib.cli.pipelines import pipeline_subcommands
+
+        pipeline_cmds = pipeline_subcommands()
+        if pipeline_cmds:
+            for subcommand, value in pipeline_cmds.items():
+                if selected == subcommand:
+                    from anomalib.cli.pipelines import PIPELINE_REGISTRY
+
+                    if PIPELINE_REGISTRY is not None:
+                        sub_parser = PIPELINE_REGISTRY[subcommand].get_parser()
+                    else:
+                        sub_parser = self.init_parser(**kwargs)
+                else:
+                    sub_parser = self.init_parser(**kwargs)
                 self.subcommand_parsers[subcommand] = sub_parser
                 parser_subcommands.add_subcommand(subcommand, sub_parser, help=value["description"])
 
@@ -179,6 +208,11 @@ def add_arguments_to_parser(parser: ArgumentParser) -> None:
 
     def add_trainer_arguments(self, parser: ArgumentParser, subcommand: str) -> None:
         """Add train arguments to the parser."""
+        from lightning.pytorch import Trainer
+
+        from anomalib.data import AnomalibDataModule
+        from anomalib.models import AnomalibModule
+
         self._add_default_arguments_to_parser(parser)
         self._add_trainer_arguments_to_parser(parser, add_optimizer=True, add_scheduler=True)
         parser.add_subclass_arguments(
@@ -199,6 +233,10 @@ def add_trainer_arguments(self, parser: ArgumentParser, subcommand: str) -> None
 
     def add_train_arguments(self, parser: ArgumentParser) -> None:
         """Add train arguments to the parser."""
+        from anomalib.data import AnomalibDataModule
+        from anomalib.engine import Engine
+        from anomalib.models import AnomalibModule
+
         self._add_default_arguments_to_parser(parser)
         self._add_trainer_arguments_to_parser(parser, add_optimizer=True, add_scheduler=True)
         parser.add_subclass_arguments(
@@ -218,6 +256,12 @@ def add_train_arguments(self, parser: ArgumentParser) -> None:
 
     def add_predict_arguments(self, parser: ArgumentParser) -> None:
         """Add predict arguments to the parser."""
+        from torch.utils.data import DataLoader, Dataset
+
+        from anomalib.data import AnomalibDataModule
+        from anomalib.engine import Engine
+        from anomalib.models import AnomalibModule
+
         self._add_default_arguments_to_parser(parser)
         self._add_trainer_arguments_to_parser(parser)
         parser.add_subclass_arguments(
@@ -241,6 +285,11 @@ def add_predict_arguments(self, parser: ArgumentParser) -> None:
 
     def add_export_arguments(self, parser: ArgumentParser) -> None:
         """Add export arguments to the parser."""
+        from anomalib.cli.utils.openvino import add_openvino_export_arguments
+        from anomalib.data import AnomalibDataModule
+        from anomalib.engine import Engine
+        from anomalib.models import AnomalibModule
+
         self._add_default_arguments_to_parser(parser)
         self._add_trainer_arguments_to_parser(parser)
         parser.add_subclass_arguments(
@@ -288,6 +337,8 @@ def _set_install_subcommand(self, action_subcommand: _ActionSubCommands) -> None
 
     def before_instantiate_classes(self) -> None:
         """Modify the configuration to properly instantiate classes and sets up tiler."""
+        from anomalib.utils.config import update_config
+
         subcommand = self.config["subcommand"]
         if subcommand in {*self.subcommands(), "train", "predict"}:
             self.config[subcommand] = update_config(self.config[subcommand])
@@ -299,6 +350,8 @@ def instantiate_classes(self) -> None:
         But for subcommands we do not want to instantiate any trainer specific classes such as datamodule, model, etc
         This is because the subcommand is responsible for instantiating and executing code based on the passed config
         """
+        from torch.utils.data import DataLoader, Dataset
+
         if self.config["subcommand"] in {*self.subcommands(), "predict"}:  # trainer commands
             # since all classes are instantiated, the LightningCLI also creates an unused ``Trainer`` object.
             # the minor change here is that engine is instantiated instead of trainer
@@ -334,6 +387,7 @@ def instantiate_engine(self) -> None:
         from lightning.pytorch.cli import SaveConfigCallback
 
         from anomalib.callbacks import get_callbacks
+        from anomalib.engine import Engine
 
         engine_args: dict[str, Any] = {}
         trainer_config = {**self._get(self.config_init, "trainer", default={}), **engine_args}
@@ -368,11 +422,14 @@ def _run_subcommand(self) -> None:
             fn = getattr(self.engine, self.subcommand)
             fn_kwargs = self._prepare_subcommand_kwargs(self.subcommand)
             fn(**fn_kwargs)
-        elif PIPELINE_REGISTRY is not None and self.subcommand in pipeline_subcommands():
-            run_pipeline(self.config)
         else:
-            self.config_init = self.parser.instantiate_classes(self.config)
-            getattr(self, f"{self.subcommand}")()
+            from anomalib.cli.pipelines import PIPELINE_REGISTRY, pipeline_subcommands, run_pipeline
+
+            if PIPELINE_REGISTRY is not None and self.subcommand in pipeline_subcommands():
+                run_pipeline(self.config)
+            else:
+                self.config_init = self.parser.instantiate_classes(self.config)
+                getattr(self, f"{self.subcommand}")()
 
     @property
     def fit(self) -> Callable:
@@ -411,6 +468,8 @@ def _add_trainer_arguments_to_parser(
         add_scheduler: bool = False,
     ) -> None:
         """Add trainer arguments to the parser."""
+        from lightning.pytorch import Trainer
+
         parser.add_class_arguments(Trainer, "trainer", fail_untyped=False, instantiate=False, sub_configs=True)
 
         if add_optimizer:
@@ -451,6 +510,10 @@ def _get(self, config: Namespace, key: str, default: Any = None) -> Any:  # noqa
 
     def _prepare_subcommand_kwargs(self, subcommand: str) -> dict[str, Any]:
         """Prepares the keyword arguments to pass to the subcommand to run."""
+        from torch.utils.data import DataLoader
+
+        from anomalib.data import AnomalibDataModule
+
         fn_kwargs = {
             k: v for k, v in self.config_init[subcommand].items() if k in self.subcommand_method_arguments[subcommand]
         }
@@ -488,6 +551,8 @@ def _configure_optimizers_method_to_model(self) -> None:
 
 def main() -> None:
     """Trainer via Anomalib CLI."""
+    from anomalib.loggers import configure_logger
+
     configure_logger()
     AnomalibCLI()
 

src/anomalib/cli/pipelines.py

@@ -7,22 +7,46 @@
 the CLI. It includes support for benchmarking and other pipeline operations.
 """
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
+
+import importlib.util
 
-from jsonargparse import Namespace
-from lightning_utilities.core.imports import module_available
+if TYPE_CHECKING:
+    from jsonargparse import Namespace
 
-from anomalib.cli.utils.help_formatter import get_short_docstring
+    from anomalib.pipelines.components.base import Pipeline
 
 logger = logging.getLogger(__name__)
 
-if module_available("anomalib.pipelines"):
-    from anomalib.pipelines import Benchmark
-    from anomalib.pipelines.components.base import Pipeline
+_UNINITIALIZED = object()
+
+_PIPELINE_REGISTRY: dict[str, type[Pipeline]] | None | object = _UNINITIALIZED
+
+_PIPELINE_DESCRIPTIONS: dict[str, str] = {
+    "benchmark": "Benchmarking pipeline for evaluating anomaly detection models.",
+}
+
+
+def _ensure_registry() -> dict[str, type[Pipeline]] | None:
+    global _PIPELINE_REGISTRY  # noqa: PLW0603
+    if _PIPELINE_REGISTRY is _UNINITIALIZED:
+        if importlib.util.find_spec("anomalib.pipelines") is not None:
+            from anomalib.pipelines import Benchmark
+
+            _PIPELINE_REGISTRY = {"benchmark": Benchmark}
+        else:
+            _PIPELINE_REGISTRY = None
+    return _PIPELINE_REGISTRY  # type: ignore[return-value]
+
 
-    PIPELINE_REGISTRY: dict[str, type[Pipeline]] | None = {"benchmark": Benchmark}
-else:
-    PIPELINE_REGISTRY = None
+def __getattr__(name: str) -> object:
+    if name == "PIPELINE_REGISTRY":
+        return _ensure_registry()
+    msg = f"module {__name__!r} has no attribute {name!r}"
+    raise AttributeError(msg)
 
 
 def pipeline_subcommands() -> dict[str, dict[str, str]]:
@@ -41,9 +65,9 @@ def pipeline_subcommands() -> dict[str, dict[str, str]]:
             }
         }
     """
-    if PIPELINE_REGISTRY is not None:
-        return {name: {"description": get_short_docstring(pipeline)} for name, pipeline in PIPELINE_REGISTRY.items()}
-    return {}
+    if importlib.util.find_spec("anomalib.pipelines") is None:
+        return {}
+    return {name: {"description": desc} for name, desc in _PIPELINE_DESCRIPTIONS.items()}
 
 
 def run_pipeline(args: Namespace) -> None:
@@ -60,10 +84,11 @@ def run_pipeline(args: Namespace) -> None:
         This feature is experimental and may change or be removed in future versions.
     """
     logger.warning("This feature is experimental. It may change or be removed in the future.")
-    if PIPELINE_REGISTRY is not None:
+    registry = _ensure_registry()
+    if registry is not None:
         subcommand = args.subcommand
         config = args[subcommand]
-        PIPELINE_REGISTRY[subcommand]().run(config)
+        registry[subcommand]().run(config)
     else:
         msg = "Pipeline is not available"
         raise ValueError(msg)