ADD import sanity checks (#226)

* ADD import sanity checks

* FIX linter

* Fix linting

* FIX json

.github/workflows/lint.yml

@@ -51,6 +51,7 @@ jobs:
           VALIDATE_MARKDOWN_PRETTIER: false
           VALIDATE_YAML_PRETTIER: false
           VALIDATE_PYTHON_PYINK: false
+          VALIDATE_JSON_PRETTIER: false
 
           # Only check new or edited files
           VALIDATE_ALL_CODEBASE: false

.vscode/extensions.json

@@ -10,6 +10,7 @@
         "ms-python.autopep8",
         "bierner.markdown-mermaid",
         "github.vscode-github-actions",
-        "tamasfe.even-better-toml"
+        "tamasfe.even-better-toml",
+        "ms-python.isort"
     ]
 }

src/itwinai/cli.py

@@ -19,6 +19,39 @@
 app = typer.Typer(pretty_exceptions_enable=False)
 
 
+@app.command()
+def sanity_check(
+    torch: Annotated[Optional[bool], typer.Option(
+        help=("Check also itwinai.torch modules.")
+    )] = False,
+    tensorflow: Annotated[Optional[bool], typer.Option(
+        help=("Check also itwinai.tensorflow modules.")
+    )] = False,
+    all: Annotated[Optional[bool], typer.Option(
+        help=("Check all modules.")
+    )] = False,
+):
+    """Run sanity checks on the installation of itwinai and
+    its dependencies by trying to import itwinai modules.
+    By default, only itwinai core modules (neither torch, nor
+    tensorflow) are tested."""
+    from itwinai.tests.sanity_check import (
+        sanity_check_all,
+        sanity_check_slim,
+        sanity_check_tensorflow,
+        sanity_check_torch,
+    )
+    all = (torch and tensorflow) or all
+    if all:
+        sanity_check_all()
+    elif torch:
+        sanity_check_torch()
+    elif tensorflow:
+        sanity_check_tensorflow()
+    else:
+        sanity_check_slim()
+
+
 @app.command()
 def scalability_report(
     pattern: Annotated[str, typer.Option(

src/itwinai/tests/__init__.py

@@ -1,7 +1,14 @@
 from .dummy_components import (
-    FakeGetter, FakeGetterExec, FakePreproc, FakePreprocExec,
-    FakeSaver, FakeSaverExec, FakeSplitter, FakeSplitterExec,
-    FakeTrainer, FakeTrainerExec
+    FakeGetter,
+    FakeGetterExec,
+    FakePreproc,
+    FakePreprocExec,
+    FakeSaver,
+    FakeSaverExec,
+    FakeSplitter,
+    FakeSplitterExec,
+    FakeTrainer,
+    FakeTrainerExec,
 )
 
 _ = (

src/itwinai/tests/dummy_components.py

@@ -1,4 +1,5 @@
 from typing import Optional
+
 from ..components import BaseComponent, monitor_exec
 
 

src/itwinai/tests/exceptions.py

@@ -0,0 +1,5 @@
+"""Custom exceptions raised during sanity checks for itwinai."""
+
+
+class SanityCheckError(Exception):
+    """Base exception for all sanity check errors."""

src/itwinai/tests/sanity_check.py

@@ -0,0 +1,104 @@
+
+import importlib
+from typing import List
+
+from .exceptions import SanityCheckError
+
+core_modules = [
+    'itwinai',
+    'itwinai.cli',
+    'itwinai.components',
+    'itwinai.distributed',
+    'itwinai.loggers',
+    'itwinai.parser',
+    'itwinai.pipeline',
+    'itwinai.serialization',
+    'itwinai.type',
+    'itwinai.utils',
+    'itwinai.tests',
+    'itwinai.tests.dummy_components',
+    'itwinai.tests.exceptions',
+    'itwinai.tests.sanity_check',
+    'itwinai.plugins'
+]
+
+torch_modules = [
+    'itwinai.torch',
+    'itwinai.torch.data',
+    'itwinai.torch.models',
+    'itwinai.torch.models.mnist',
+    'itwinai.torch.config',
+    'itwinai.torch.distributed',
+    'itwinai.torch.inference',
+    'itwinai.torch.mlflow',
+    'itwinai.torch.reproducibility',
+    'itwinai.torch.trainer',
+    'itwinai.torch.type',
+]
+
+tensorflow_modules = [
+    'itwinai.tensorflow',
+    'itwinai.tensorflow.data',
+    'itwinai.tensorflow.models',
+    'itwinai.tensorflow.models.mnist',
+    'itwinai.tensorflow.distributed',
+    'itwinai.tensorflow.trainer',
+    'itwinai.tensorflow.utils',
+]
+
+
+def run_sanity_check(modules: List[str]):
+    """Run sanity checks by trying to import modules.
+
+    Args:
+        modules (List[str]): list of modules
+
+    Raises:
+        SanityCheckError: when some module cannot be imported.
+    """
+    failed_imports = []
+
+    for module in modules:
+        try:
+            importlib.import_module(module)
+            print(f"✅ Successfully imported: {module}")
+        except ImportError as e:
+            failed_imports.append((module, str(e)))
+            print(f"❌ Failed to import: {module} - {e}")
+
+    if failed_imports:
+        err_msg = "\nSummary of failed imports:\n"
+        for module, error in failed_imports:
+            err_msg += f"Module: {module}, Error: {error}\n"
+
+        raise SanityCheckError(
+            "Not all itwinai modules could be successfully imported!\n"
+            + err_msg
+        )
+    else:
+        print("\nAll modules imported successfully!")
+
+
+def sanity_check_slim():
+    """Run sanity check on the installation
+    of core modules of itwinai (neither itwinai.torch,
+    nor itwinai.tensorflow)."""
+
+    run_sanity_check(modules=core_modules)
+
+
+def sanity_check_torch():
+    """Run sanity check on the installation of itwinai
+    for a torch environment."""
+    run_sanity_check(modules=core_modules+torch_modules)
+
+
+def sanity_check_tensorflow():
+    """Run sanity check on the installation of itwinai
+    for a tensorflow environment."""
+    run_sanity_check(modules=core_modules+tensorflow_modules)
+
+
+def sanity_check_all():
+    """Run all sanity checks."""
+    run_sanity_check(modules=core_modules+torch_modules+tensorflow_modules)