Support multiple backends

.pre-commit-config.yaml

@@ -15,6 +15,11 @@ repos:
       - id: check-builtin-literals
       - id: trailing-whitespace
 
+  - repo: https://github.com/PyCQA/autoflake
+    rev: v2.3.1
+    hooks:
+      - id: autoflake
+
   - repo: https://github.com/psf/black
     rev: 24.2.0
     hooks:

conftest.py

@@ -1,16 +1,33 @@
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
 from jupyter_scheduler.orm import Base
 from jupyter_scheduler.scheduler import Scheduler
-from jupyter_scheduler.tests.mocks import MockEnvironmentManager
+from jupyter_scheduler.tests.mocks import MockEnvironmentManager, MockTestBackend
 
 pytest_plugins = ("jupyter_server.pytest_plugin", "pytest_jupyter.jupyter_server")
 
 
+def _mock_discover_backends(*args, **kwargs):
+    """Return test backends for testing."""
+    from jupyter_scheduler.backends import JupyterServerNotebookBackend
+
+    return {"jupyter_server_nb": JupyterServerNotebookBackend, "test": MockTestBackend}
+
+
+@pytest.fixture(autouse=True)
+def mock_backend_discovery():
+    """Patch backend discovery to include test backend for all tests."""
+    with patch(
+        "jupyter_scheduler.extension.discover_backends", side_effect=_mock_discover_backends
+    ):
+        yield
+
+
 @pytest.fixture(scope="session")
 def static_test_files_dir() -> Path:
     return Path(__file__).parent.resolve() / "jupyter_scheduler" / "tests" / "static"
@@ -51,6 +68,7 @@ def jp_scheduler_db(jp_scheduler_db_url):
     session = Session()
     yield session
     session.close()
+    engine.dispose()
 
 
 @pytest.fixture

dev/seed.py

@@ -155,7 +155,7 @@ async def load_data(jobs_count: int, job_defs_count: int, db_path: str):
         f"\nCreated {jobs_count} jobs and {job_defs_count} job definitions in the scheduler database"
     )
     click.echo(f"present at {db_path}. Copy the following command")
-    click.echo(f"to start JupyterLab with this database.\n")
+    click.echo("to start JupyterLab with this database.\n")
     click.echo(f"`jupyter lab --SchedulerApp.db_url={db_url}`\n")
 
 

jupyter_scheduler/__init__.py

@@ -1,6 +1,5 @@
 """Scheduling API for JupyterLab"""
 
-from ._version import __version__
 from .extension import SchedulerApp
 
 

jupyter_scheduler/_version.py

@@ -1,6 +1,3 @@
-import json
-from pathlib import Path
-
 __all__ = ["__version__"]
 
 version_info = (2, 11, 0, "", "")

jupyter_scheduler/backend_registry.py

@@ -0,0 +1,167 @@
+import logging
+from typing import Any, Dict, List, Optional, Type
+
+from jupyter_scheduler.backends import BackendConfig, DescribeBackendResponse
+from jupyter_scheduler.environments import EnvironmentManager
+from jupyter_scheduler.orm import create_tables
+from jupyter_scheduler.pydantic_v1 import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+def import_class(class_path: str) -> Type:
+    """Import a class from a fully qualified path like 'module.submodule.ClassName'."""
+    module_path, class_name = class_path.rsplit(".", 1)
+    module = __import__(module_path, fromlist=[class_name])
+    return getattr(module, class_name)
+
+
+class BackendInstance(BaseModel):
+    """A running backend with its configuration and initialized scheduler."""
+
+    config: BackendConfig
+    scheduler: Any  # BaseScheduler at runtime, but Any to support test mocks
+
+
+class BackendRegistry:
+    """Registry for storing, initializing, and routing to scheduler backends."""
+
+    def __init__(
+        self,
+        configs: List[BackendConfig],
+        legacy_job_backend: str,
+        preferred_backends: Optional[Dict[str, str]] = None,
+    ):
+        self._configs = configs
+        self._backends: Dict[str, BackendInstance] = {}
+        self._legacy_job_backend = legacy_job_backend
+        self._preferred_backends = preferred_backends or {}
+        self._extension_map: Dict[str, List[str]] = {}
+
+    def initialize(
+        self,
+        root_dir: str,
+        environments_manager: EnvironmentManager,
+        db_url: str,
+        config: Optional[Any] = None,
+    ):
+        """Instantiate all backends from configs."""
+        seen_ids = set()
+        for cfg in self._configs:
+            if cfg.id in seen_ids:
+                raise ValueError(f"Duplicate backend ID: '{cfg.id}'")
+            if ":" in cfg.id:
+                raise ValueError(f"Backend ID cannot contain ':': '{cfg.id}'")
+            seen_ids.add(cfg.id)
+
+        for cfg in self._configs:
+            try:
+                instance = self._create_backend(cfg, root_dir, environments_manager, db_url, config)
+                self._backends[cfg.id] = instance
+
+                for ext in cfg.file_extensions:
+                    ext_lower = ext.lower().lstrip(".")
+                    if ext_lower not in self._extension_map:
+                        self._extension_map[ext_lower] = []
+                    self._extension_map[ext_lower].append(cfg.id)
+
+                logger.info(f"Initialized backend: {cfg.id} ({cfg.name})")
+            except Exception as e:
+                logger.error(f"Failed to initialize backend {cfg.id}: {e}")
+                raise
+
+    def _create_backend(
+        self,
+        cfg: BackendConfig,
+        root_dir: str,
+        environments_manager: EnvironmentManager,
+        global_db_url: str,
+        config: Optional[Any] = None,
+    ) -> BackendInstance:
+        """Import scheduler class, instantiate it, and return a BackendInstance.
+
+        Creates database tables if not found and backend uses default SQLAlchemy storage.
+        """
+        scheduler_class = import_class(cfg.scheduler_class)
+
+        backend_db_url = cfg.db_url or global_db_url
+
+        # Create SQL tables only if backend uses default SQLAlchemy storage.
+        # Backends with custom database_manager_class handle their own storage.
+        if backend_db_url and cfg.database_manager_class is None:
+            create_tables(backend_db_url)
+
+        scheduler = scheduler_class(
+            root_dir=root_dir,
+            environments_manager=environments_manager,
+            db_url=backend_db_url,
+            config=config,
+            backend_id=cfg.id,
+        )
+
+        if cfg.execution_manager_class:
+            scheduler.execution_manager_class = import_class(cfg.execution_manager_class)
+
+        return BackendInstance(config=cfg, scheduler=scheduler)
+
+    def get_backend(self, backend_id: str) -> Optional[BackendInstance]:
+        """Return a backend with matching ID, None if none is found."""
+        return self._backends.get(backend_id)
+
+    def get_legacy_job_backend(self) -> BackendInstance:
+        """Get the backend for routing legacy jobs (UUID-only IDs from pre-3.0).
+
+        Raises:
+            KeyError: If the configured legacy_job_backend ID is not found.
+        """
+        if self._legacy_job_backend not in self._backends:
+            raise KeyError(f"Legacy job backend '{self._legacy_job_backend}' not found in registry")
+        return self._backends[self._legacy_job_backend]
+
+    def get_for_file(self, input_uri: str) -> BackendInstance:
+        """Auto-select backend by file extension. Prefers configured backend, else alphabetical.
+
+        Raises:
+            ValueError: If no backend supports the file extension.
+        """
+        ext = ""
+        if "." in input_uri:
+            ext = input_uri.rsplit(".", 1)[-1].lower()
+
+        candidate_ids = self._extension_map.get(ext, [])
+        if not candidate_ids:
+            raise ValueError(f"No backend supports file extension '.{ext}'")
+
+        # 1. Check explicit preference for this extension
+        preferred_id = self._preferred_backends.get(ext)
+        if preferred_id and preferred_id in candidate_ids:
+            return self._backends[preferred_id]
+
+        # 2. Otherwise return min by name (first alphabetically)
+        candidate_instances = [self._backends[bid] for bid in candidate_ids]
+        return min(candidate_instances, key=lambda b: b.config.name)
+
+    def describe_backends(self) -> List[DescribeBackendResponse]:
+        """Return backend descriptions sorted alphabetically by name. Frontend uses first as default."""
+        backends_sorted = sorted(self._backends.values(), key=lambda b: b.config.name)
+        return [
+            DescribeBackendResponse(
+                id=b.config.id,
+                name=b.config.name,
+                description=b.config.description,
+                file_extensions=b.config.file_extensions,
+                output_formats=b.config.output_formats,
+            )
+            for b in backends_sorted
+        ]
+
+    @property
+    def backends(self) -> List[BackendInstance]:
+        """Return all backend instances."""
+        return list(self._backends.values())
+
+    def __len__(self) -> int:
+        return len(self._backends)
+
+    def __contains__(self, backend_id: str) -> bool:
+        return backend_id in self._backends

jupyter_scheduler/backend_utils.py

@@ -0,0 +1,70 @@
+import logging
+from importlib.metadata import entry_points
+from typing import Dict, Optional, Type
+
+from jupyter_scheduler.backends import DEFAULT_FALLBACK_BACKEND_ID
+from jupyter_scheduler.base_backend import BaseBackend
+
+ENTRY_POINT_GROUP = "jupyter_scheduler.backends"
+
+logger = logging.getLogger(__name__)
+
+
+def discover_backends(
+    log: Optional[logging.Logger] = None,
+) -> Dict[str, Type[BaseBackend]]:
+    """Discover backends registered in the 'jupyter_scheduler.backends' entry point group."""
+    if log is None:
+        log = logger
+
+    backends: Dict[str, Type[BaseBackend]] = {}
+
+    eps = entry_points()
+    if hasattr(eps, "select"):
+        backend_eps = eps.select(group=ENTRY_POINT_GROUP)
+    else:
+        backend_eps = eps.get(ENTRY_POINT_GROUP, [])
+
+    for ep in backend_eps:
+        try:
+            backend_class = ep.load()
+        except ImportError as e:
+            missing_package = getattr(e, "name", str(e))
+            log.warning(
+                f"Unable to load backend '{ep.name}': missing dependency '{missing_package}'. "
+                f"Install the required package to enable this backend."
+            )
+            continue
+        except Exception as e:
+            log.warning(f"Unable to load backend '{ep.name}': {e}")
+            continue
+
+        if not hasattr(backend_class, "id"):
+            log.warning(f"Backend '{ep.name}' does not define 'id' attribute. Skipping.")
+            continue
+
+        backend_id = backend_class.id
+        backends[backend_id] = backend_class
+        log.info(f"Registered backend '{backend_id}' ({backend_class.name})")
+
+    return backends
+
+
+def get_legacy_job_backend_id(
+    available_backends: Dict[str, Type[BaseBackend]],
+    legacy_job_backend: Optional[str] = None,
+) -> str:
+    """Get backend ID for routing legacy jobs (UUID-only IDs from pre-3.0)."""
+    if not available_backends:
+        raise ValueError("No scheduler backends available.")
+
+    if legacy_job_backend and legacy_job_backend in available_backends:
+        return legacy_job_backend
+
+    if DEFAULT_FALLBACK_BACKEND_ID in available_backends:
+        return DEFAULT_FALLBACK_BACKEND_ID
+
+    raise ValueError(
+        f"No backend for legacy jobs. Set SchedulerApp.legacy_job_backend. "
+        f"Available: {list(available_backends.keys())}"
+    )

jupyter_scheduler/backends.py

@@ -0,0 +1,72 @@
+from typing import Any, Dict, List, Optional
+
+from jupyter_scheduler.base_backend import BaseBackend
+from jupyter_scheduler.models import OutputFormat
+from jupyter_scheduler.pydantic_v1 import BaseModel, Field
+
+JUPYTER_SERVER_NB_BACKEND_ID = "jupyter_server_nb"
+JUPYTER_SERVER_PY_BACKEND_ID = "jupyter_server_py"
+DEFAULT_FALLBACK_BACKEND_ID = JUPYTER_SERVER_NB_BACKEND_ID
+
+
+class BackendConfig(BaseModel):
+    """Runtime configuration for an initialized backend instance."""
+
+    id: str
+    name: str
+    description: str
+    scheduler_class: str
+    execution_manager_class: str
+    database_manager_class: Optional[str] = None
+    db_url: Optional[str] = None
+    file_extensions: List[str] = Field(default_factory=list)
+    output_formats: List[Dict[str, str]] = Field(default_factory=list)
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class DescribeBackendResponse(BaseModel):
+    """API response model for GET /scheduler/backends.
+
+    Backends are returned sorted alphabetically by name for consistent UI ordering.
+    Use preferred_backends config to control which backend is pre-selected per file extension.
+    """
+
+    id: str
+    name: str
+    description: str
+    file_extensions: List[str]
+    output_formats: List[OutputFormat]
+
+    class Config:
+        orm_mode = True
+
+
+class JupyterServerNotebookBackend(BaseBackend):
+    """Built-in backend executing notebooks via nbconvert on the Jupyter server."""
+
+    id = JUPYTER_SERVER_NB_BACKEND_ID
+    name = "Jupyter Server Notebook"
+    description = "Execute notebooks on the Jupyter server"
+    scheduler_class = "jupyter_scheduler.scheduler.Scheduler"
+    execution_manager_class = "jupyter_scheduler.executors.DefaultExecutionManager"
+    file_extensions = ["ipynb"]
+    output_formats = [
+        {"id": "ipynb", "label": "Notebook", "description": "Executed notebook with outputs"},
+        {"id": "html", "label": "HTML", "description": "HTML export of notebook"},
+    ]
+
+
+class JupyterServerPythonBackend(BaseBackend):
+    """Built-in backend executing Python scripts via subprocess on the Jupyter server."""
+
+    id = JUPYTER_SERVER_PY_BACKEND_ID
+    name = "Jupyter Server Python"
+    description = "Execute Python scripts on the Jupyter server"
+    scheduler_class = "jupyter_scheduler.scheduler.Scheduler"
+    execution_manager_class = "jupyter_scheduler.python_executor.PythonScriptExecutionManager"
+    file_extensions = ["py"]
+    output_formats = [
+        {"id": "stdout", "label": "Output", "description": "Standard output from script"},
+        {"id": "stderr", "label": "Errors", "description": "Standard error from script"},
+        {"id": "json", "label": "JSON", "description": "JSON result if script produces one"},
+    ]