feat: use server_lifespan

feat: use `server_lifespan`

Author: cofin

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ default_language_version:
   python: "3.11"
 repos:
   - repo: https://github.com/compilerla/conventional-pre-commit
-    rev: v2.4.0
+    rev: v3.0.0
     hooks:
       - id: conventional-pre-commit
         stages: [commit-msg]
@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.1.0"
+    rev: "v0.1.8"
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -28,12 +28,12 @@ repos:
       - id: codespell
         exclude: "pdm.lock"
   - repo: https://github.com/psf/black
-    rev: 23.10.0
+    rev: 23.12.0
     hooks:
       - id: black
         args: [--config=./pyproject.toml]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.6.0"
+    rev: "v1.7.1"
     hooks:
       - id: mypy
         exclude: "docs"
@@ -53,6 +53,6 @@ repos:
             "litestar[cli]",
           ]
   - repo: https://github.com/sphinx-contrib/sphinx-lint
-    rev: "v0.8.1"
+    rev: "v0.9.1"
     hooks:
       - id: sphinx-lint

examples/basic.py

@@ -26,6 +26,7 @@ async def samples_queue_info(self, task_queues: TaskQueues) -> QueueInfo:
     config=SAQConfig(
         redis_url="redis://localhost:6397/0",
         web_enabled=True,
+        use_server_lifespan=True,
         queue_configs=[
             QueueConfig(
                 name="samples",

litestar_saq/_util.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-import pkgutil
+import importlib.util
 import sys
 from functools import lru_cache
 from importlib import import_module
-from importlib.machinery import SourceFileLoader
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -24,11 +23,11 @@ def module_to_os_path(dotted_path: str) -> Path:
 
     Ensures that pkgutil returns a valid source file loader.
     """
-    src = pkgutil.get_loader(dotted_path)
-    if not isinstance(src, SourceFileLoader):
+    src = importlib.util.find_spec(dotted_path)
+    if src is None:
         msg = f"Couldn't find the path for {dotted_path}"
         raise TypeError(msg)
-    return Path(str(src.path).removesuffix("/__init__.py"))
+    return Path(str(src.origin).removesuffix("/__init__.py"))  # type: ignore[unreachable]
 
 
 def import_string(dotted_path: str) -> Any:

litestar_saq/base.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from saq import Job as SaqJob
 from saq import Worker as SaqWorker
@@ -106,8 +106,10 @@ def __init__(
         timers: PartialTimersDict | None = None,
         dequeue_timeout: float = 0,
         separate_process: bool = True,
+        multiprocessing_mode: Literal["multiprocessing", "threading"] = "multiprocessing",
     ) -> None:
         self.separate_process = separate_process
+        self.multiprocessing_mode = multiprocessing_mode
         super().__init__(
             cast("SaqQueue", queue),
             functions,

litestar_saq/cli.py

@@ -53,11 +53,11 @@ def run_worker(
         show_saq_info(app, workers, plugin)
         if workers > 1:
             for _ in range(workers - 1):
-                p = multiprocessing.Process(target=run_worker_process, args=(plugin.get_workers(), app.logging_config))
+                p = multiprocessing.Process(target=run_saq_worker, args=(plugin.get_workers(), app.logging_config))
                 p.start()
 
         try:
-            run_worker_process(
+            run_saq_worker(
                 workers=plugin.get_workers(),
                 logging_config=cast("BaseLoggingConfig", app.logging_config),
             )
@@ -84,7 +84,7 @@ def worker_status(
         if debug is not None or verbose is not None:
             app.debug = True
         plugin = get_saq_plugin(app)
-        show_saq_info(app, 1, plugin)
+        show_saq_info(app, plugin.config.worker_processes, plugin)
 
     return background_worker_group
 
@@ -128,7 +128,7 @@ def show_saq_info(app: Litestar, workers: int, plugin: SAQPlugin) -> None:  # pr
     console.print(table)
 
 
-def run_worker_process(workers: list[Worker], logging_config: BaseLoggingConfig | None) -> None:
+def run_saq_worker(workers: list[Worker], logging_config: BaseLoggingConfig | None) -> None:
     """Run a worker."""
     import asyncio
 

litestar_saq/config.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Collection, Dict, Mapping, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Callable, Collection, Dict, Literal, Mapping, TypeVar, Union, cast
 
 from litestar.exceptions import ImproperlyConfiguredException
 from litestar.serialization import decode_json, encode_json
@@ -40,7 +40,7 @@ def serializer(value: Any) -> str:
 
 
 def _get_static_files() -> Path:
-    return Path(module_to_os_path("saq.web") / "static")
+    return Path(module_to_os_path("saq") / "web" / "static")
 
 
 @dataclass
@@ -80,8 +80,7 @@ class SAQConfig:
 
     Default is set to 1.
     """
-    web_enabled: bool = False
-    """If true, the worker admin UI is launched on worker startup.."""
+
     json_deserializer: LoadType = decode_json
     """This is a Python callable that will
     convert a JSON string to a Python object. By default, this is set to Litestar's
@@ -91,10 +90,16 @@ class SAQConfig:
     By default, Litestar's :attr:`encode_json() <.serialization.encode_json>` is used."""
     static_files: Path = field(default_factory=_get_static_files)
     """Location of the static files to serve for the SAQ UI"""
+    web_enabled: bool = False
+    """If true, the worker admin UI is launched on worker startup.."""
     web_path = "/saq"
     """Base path to serve the SAQ web UI"""
     web_guards: list[Guard] | None = field(default=None)
     """Guards to apply to web endpoints."""
+    web_include_in_schema: bool = True
+    """Include Queue API endpoints in generated OpenAPI schema"""
+    use_server_lifespan: bool = False
+    """Utilize the server lifespan hook to run SAQ."""
 
     def __post_init__(self) -> None:
         if self.redis is not None and self.redis_url is not None:
@@ -189,7 +194,7 @@ class QueueConfig:
     concurrency: int = 10
     """Number of jobs to process concurrently"""
     max_concurrent_ops: int = 15
-    """Maximum concurrent operations. (default 20)
+    """Maximum concurrent operations. (default 15)
             This throttles calls to `enqueue`, `job`, and `abort` to prevent the Queue
             from consuming too many Redis connections."""
     tasks: Collection[ReceivesContext | tuple[str, Function] | str] = field(default_factory=list)
@@ -212,6 +217,9 @@ class QueueConfig:
             abort: how often to check if a job is aborted"""
     dequeue_timeout: float = 0
     """How long it will wait to dequeue"""
+    multiprocessing_mode: Literal["multiprocessing", "threading"] = "multiprocessing"
+    """Executes with the multiprocessing or threading backend.
+            Set it threading for workloads that aren't CPU bound."""
     separate_process: bool = True
     """Executes as a separate event loop when True.
             Set it False to execute within the Litestar application."""

litestar_saq/plugin.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Collection, TypeVar, cast
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Collection, Iterator, TypeVar, cast
 
-from litestar.di import Provide
-from litestar.plugins import CLIPluginProtocol, InitPluginProtocol
-from litestar.static_files import StaticFilesConfig
+from litestar.plugins import CLIPlugin, InitPluginProtocol
 
 from litestar_saq.base import Worker
-from litestar_saq.controllers import build_controller
 
 if TYPE_CHECKING:
     from click import Group
+    from litestar import Litestar
     from litestar.config.app import AppConfig
     from saq.types import Function
 
@@ -20,7 +19,7 @@
 T = TypeVar("T")
 
 
-class SAQPlugin(InitPluginProtocol, CLIPluginProtocol):
+class SAQPlugin(InitPluginProtocol, CLIPlugin):
     """SAQ plugin."""
 
     __slots__ = ("_config", "_worker_instances")
@@ -34,6 +33,10 @@ def __init__(self, config: SAQConfig) -> None:
         self._config = config
         self._worker_instances: list[Worker] | None = None
 
+    @property
+    def config(self) -> SAQConfig:
+        return self._config
+
     def on_cli_init(self, cli: Group) -> None:
         from litestar_saq.cli import build_cli_app
 
@@ -46,6 +49,12 @@ def on_app_init(self, app_config: AppConfig) -> AppConfig:
         Args:
             app_config: The :class:`AppConfig <.config.app.AppConfig>` instance.
         """
+
+        from litestar.di import Provide
+        from litestar.static_files import StaticFilesConfig
+
+        from litestar_saq.controllers import build_controller
+
         app_config.dependencies.update(
             {
                 self._config.queues_dependency_key: Provide(
@@ -102,3 +111,27 @@ def get_queues(self) -> TaskQueues:
 
     def get_queue(self, name: str) -> Queue:
         return self.get_queues().get(name)
+
+    @contextmanager
+    def server_lifespan(self, app: Litestar) -> Iterator[None]:
+        import multiprocessing
+
+        from litestar_saq.cli import run_saq_worker
+
+        if self._config.use_server_lifespan:
+            processes = [
+                multiprocessing.Process(target=run_saq_worker, args=(self.get_workers(), app.logging_config))
+                for _ in range(self._config.worker_processes)
+            ]
+
+            try:
+                for p in processes:
+                    p.start()
+                yield
+            finally:
+                for p in processes:
+                    if p.is_alive():
+                        p.terminate()
+                        p.join()
+        else:
+            yield