beamflow-runtime 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- beamflow_runtime-0.3.0/PKG-INFO +18 -0
- beamflow_runtime-0.3.0/beamflow_runtime/__init__.py +39 -0
- beamflow_runtime-0.3.0/beamflow_runtime/ingress/__init__.py +2 -0
- beamflow_runtime-0.3.0/beamflow_runtime/ingress/app.py +74 -0
- beamflow_runtime-0.3.0/beamflow_runtime/ingress/router.py +68 -0
- beamflow_runtime-0.3.0/beamflow_runtime/wiring/__init__.py +2 -0
- beamflow_runtime-0.3.0/beamflow_runtime/wiring/runtime.py +161 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/__init__.py +2 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/__init__.py +0 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/__init__.py +0 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/bootstrap.py +57 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/consumer_middleware.py +48 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/dramatiq_backend.py +380 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/dramatiq_consumer.py +77 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/dramatiq/entrypoint.py +30 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/managed/__init__.py +0 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/managed/http_entrypoint.py +208 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/managed/main.py +20 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/backends/managed/managed_consumer.py +27 -0
- beamflow_runtime-0.3.0/beamflow_runtime/worker/runner.py +89 -0
- beamflow_runtime-0.3.0/pyproject.toml +28 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: beamflow-runtime
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Actual runtime engine and bootstrap layer for the Beamflow platform.
|
|
5
|
+
Author: juraj.bezdek@gmail.com
|
|
6
|
+
Author-email: juraj.bezdek@gmail.com
|
|
7
|
+
Requires-Python: >=3.11
|
|
8
|
+
Classifier: Programming Language :: Python :: 3
|
|
9
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
10
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
12
|
+
Requires-Dist: apscheduler (>=3.11.0)
|
|
13
|
+
Requires-Dist: beamflow-clients (>=0.3.0,<0.4.0)
|
|
14
|
+
Requires-Dist: beamflow-lib (>=0.3.0,<0.4.0)
|
|
15
|
+
Requires-Dist: dramatiq (>=1.16.0)
|
|
16
|
+
Requires-Dist: fastapi (>=0.110.0)
|
|
17
|
+
Requires-Dist: setuptools (<70.0.0)
|
|
18
|
+
Requires-Dist: uvicorn (>=0.29.0)
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Framework Runtime - Pure glue around beamflow_lib.
|
|
3
|
+
|
|
4
|
+
This package provides:
|
|
5
|
+
- Ingress API helpers (FastAPI router generation from webhook registry)
|
|
6
|
+
- Worker bootstrap (Dramatiq broker configuration with context middleware)
|
|
7
|
+
- Queue backends (Dramatiq implementation moved from beamflow_lib)
|
|
8
|
+
|
|
9
|
+
Usage for Ingress API:
|
|
10
|
+
from beamflow_runtime.ingress.app import create_fastapi_app
|
|
11
|
+
from beamflow_lib.config.loader import load_runtime_config
|
|
12
|
+
|
|
13
|
+
config = load_runtime_config("/path/to/config")
|
|
14
|
+
app = create_fastapi_app(config)
|
|
15
|
+
|
|
16
|
+
Usage for Worker:
|
|
17
|
+
from beamflow_runtime.worker.entrypoint import configure_worker_process
|
|
18
|
+
from beamflow_lib.config.loader import load_runtime_config
|
|
19
|
+
|
|
20
|
+
config = load_runtime_config("/path/to/config")
|
|
21
|
+
runtime = configure_worker_process(config)
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
# Main exports
|
|
25
|
+
from .wiring.runtime import Runtime, initialize_runtime, build_worker_runtime
|
|
26
|
+
from .ingress.app import create_fastapi_app
|
|
27
|
+
from .ingress.router import build_webhook_router
|
|
28
|
+
from .worker.runner import run_worker, build_worker_consumer
|
|
29
|
+
|
|
30
|
+
__all__ = [
|
|
31
|
+
"Runtime",
|
|
32
|
+
"initialize_runtime",
|
|
33
|
+
|
|
34
|
+
"build_worker_runtime",
|
|
35
|
+
"create_fastapi_app",
|
|
36
|
+
"build_webhook_router",
|
|
37
|
+
"run_worker",
|
|
38
|
+
"build_worker_consumer",
|
|
39
|
+
]
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
"""
|
|
2
|
+
FastAPI application factory for webhook ingress.
|
|
3
|
+
|
|
4
|
+
Provides a convenience function to create a fully configured FastAPI app.
|
|
5
|
+
"""
|
|
6
|
+
from fastapi import FastAPI
|
|
7
|
+
from beamflow_lib.config.runtime_config import RuntimeConfig
|
|
8
|
+
from .router import build_webhook_router
|
|
9
|
+
from ..wiring.runtime import initialize_runtime
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
from typing import List, Union, Optional
|
|
13
|
+
from pathlib import Path
|
|
14
|
+
|
|
15
|
+
def create_fastapi_app(
|
|
16
|
+
config: RuntimeConfig,
|
|
17
|
+
auto_import: Optional[List[Union[str, Path]]] = None
|
|
18
|
+
) -> FastAPI:
|
|
19
|
+
"""
|
|
20
|
+
Create FastAPI app with webhook routes.
|
|
21
|
+
|
|
22
|
+
Usage in ingress service:
|
|
23
|
+
from beamflow_runtime.ingress.app import create_fastapi_app
|
|
24
|
+
from beamflow_lib.config.loader import load_runtime_config
|
|
25
|
+
|
|
26
|
+
config = load_runtime_config("/path/to/config")
|
|
27
|
+
app = create_fastapi_app(config)
|
|
28
|
+
|
|
29
|
+
Args:
|
|
30
|
+
config: RuntimeConfig with webhooks configuration
|
|
31
|
+
auto_import: List of paths to auto-import modules from
|
|
32
|
+
|
|
33
|
+
Returns:
|
|
34
|
+
FastAPI application with webhook routes included
|
|
35
|
+
"""
|
|
36
|
+
# Build runtime (imports modules, sets up backend)
|
|
37
|
+
rt = initialize_runtime(config, auto_import=auto_import)
|
|
38
|
+
|
|
39
|
+
# Create app
|
|
40
|
+
app = FastAPI(title="Ingress API")
|
|
41
|
+
|
|
42
|
+
# Include webhook router with configured prefix
|
|
43
|
+
app.include_router(build_webhook_router(), prefix=config.webhooks.prefix)
|
|
44
|
+
|
|
45
|
+
# Store runtime reference on app for access in middleware/dependencies if needed
|
|
46
|
+
app.state.runtime = rt
|
|
47
|
+
|
|
48
|
+
import logging
|
|
49
|
+
import sys
|
|
50
|
+
app_logger = logging.getLogger(__name__)
|
|
51
|
+
app_logger.setLevel(logging.INFO)
|
|
52
|
+
if not app_logger.handlers:
|
|
53
|
+
handler = logging.StreamHandler(sys.stdout)
|
|
54
|
+
handler.setFormatter(logging.Formatter('INFO: %(message)s'))
|
|
55
|
+
app_logger.addHandler(handler)
|
|
56
|
+
|
|
57
|
+
# Include auto-imported routers
|
|
58
|
+
if hasattr(rt, 'imported_modules'):
|
|
59
|
+
for mod in rt.imported_modules:
|
|
60
|
+
if hasattr(mod, "router"):
|
|
61
|
+
app.include_router(mod.router)
|
|
62
|
+
mod_file = getattr(mod, "__file__", mod.__name__)
|
|
63
|
+
app_logger.info(f"router automimpoted from [{mod_file}]")
|
|
64
|
+
|
|
65
|
+
app_logger.info("Registered paths:")
|
|
66
|
+
for route in app.routes:
|
|
67
|
+
methods = getattr(route, "methods", None)
|
|
68
|
+
path = getattr(route, "path", None)
|
|
69
|
+
name = getattr(route, "name", getattr(route, "endpoint", "").__name__ if hasattr(route, "endpoint") else "")
|
|
70
|
+
if path:
|
|
71
|
+
methods_str = ",".join(methods) if methods else "ANY"
|
|
72
|
+
app_logger.info(f"- {methods_str} {path} [{name}]")
|
|
73
|
+
|
|
74
|
+
return app
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Webhook router generation from beamflow_lib registry.
|
|
3
|
+
|
|
4
|
+
Builds FastAPI routes from webhooks registered with @ingress.webhook decorator.
|
|
5
|
+
"""
|
|
6
|
+
import inspect
|
|
7
|
+
from typing import Any, Optional
|
|
8
|
+
from fastapi import APIRouter, Request, Response
|
|
9
|
+
from beamflow_lib.pipelines.ingress import ingress
|
|
10
|
+
from beamflow_lib.context import integration_context
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
async def _invoke_handler(handler, request: Request, body: bytes) -> Any:
|
|
14
|
+
"""
|
|
15
|
+
Invoke a webhook handler within the integration context.
|
|
16
|
+
|
|
17
|
+
Handles both sync and async handlers.
|
|
18
|
+
"""
|
|
19
|
+
metadata = handler._ingress_metadata
|
|
20
|
+
|
|
21
|
+
with integration_context(
|
|
22
|
+
integration=metadata["integration"],
|
|
23
|
+
integration_pipeline=metadata["pipeline"]
|
|
24
|
+
):
|
|
25
|
+
# Pass the request to the handler - let the handler decide what to do with it
|
|
26
|
+
if inspect.iscoroutinefunction(handler):
|
|
27
|
+
result = await handler(request)
|
|
28
|
+
else:
|
|
29
|
+
result = handler(request)
|
|
30
|
+
|
|
31
|
+
return result
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def build_webhook_router() -> APIRouter:
|
|
35
|
+
"""
|
|
36
|
+
Build FastAPI router from registered webhooks.
|
|
37
|
+
|
|
38
|
+
Reads all webhooks registered via @ingress.webhook and creates
|
|
39
|
+
FastAPI endpoints for each one.
|
|
40
|
+
|
|
41
|
+
Returns:
|
|
42
|
+
APIRouter with webhook routes tagged under "webhooks"
|
|
43
|
+
"""
|
|
44
|
+
router = APIRouter(tags=["webhooks"])
|
|
45
|
+
|
|
46
|
+
for handler in ingress.get_webhooks():
|
|
47
|
+
metadata = handler._ingress_metadata
|
|
48
|
+
# Build path from metadata
|
|
49
|
+
path = metadata["path"]
|
|
50
|
+
methods = [metadata.get("method", "POST")]
|
|
51
|
+
|
|
52
|
+
# Create endpoint closure that captures the handler
|
|
53
|
+
# Define the endpoint function factory
|
|
54
|
+
def create_endpoint(h=handler):
|
|
55
|
+
async def endpoint(request: Request) -> Response:
|
|
56
|
+
result = await _invoke_handler(h, request, await request.body())
|
|
57
|
+
# Default to 202 Accepted for webhook processing
|
|
58
|
+
if result is None:
|
|
59
|
+
return Response(status_code=202)
|
|
60
|
+
return result
|
|
61
|
+
return endpoint
|
|
62
|
+
|
|
63
|
+
# Create the endpoint function
|
|
64
|
+
endpoint = create_endpoint(handler)
|
|
65
|
+
|
|
66
|
+
router.add_api_route(path, endpoint, methods=methods)
|
|
67
|
+
|
|
68
|
+
return router
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
from typing import Any, Optional, List, Union
|
|
5
|
+
import importlib
|
|
6
|
+
import importlib.util
|
|
7
|
+
import dramatiq
|
|
8
|
+
from dramatiq.brokers.redis import RedisBroker
|
|
9
|
+
from beamflow_lib.config.runtime_config import RuntimeConfig, BackendType
|
|
10
|
+
from beamflow_lib.queue.backend import set_backend, TaskBackend
|
|
11
|
+
from beamflow_runtime.worker.backends.dramatiq.dramatiq_backend import DramatiqBackend, FrameworkContextMiddleware
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass
|
|
15
|
+
class Runtime:
|
|
16
|
+
"""Runtime container for backend and config references."""
|
|
17
|
+
backend: TaskBackend
|
|
18
|
+
config: RuntimeConfig
|
|
19
|
+
imported_modules: List[Any] = field(default_factory=list)
|
|
20
|
+
|
|
21
|
+
@property
|
|
22
|
+
def broker(self):
|
|
23
|
+
"""Return the global Dramatiq broker."""
|
|
24
|
+
import dramatiq
|
|
25
|
+
return dramatiq.get_broker()
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _auto_import_path(path: Union[str, Path]) -> List[Any]:
|
|
29
|
+
"""
|
|
30
|
+
Import modules from a file or directory path.
|
|
31
|
+
|
|
32
|
+
If path is a directory:
|
|
33
|
+
- If it contains __init__.py, import it as a module.
|
|
34
|
+
- Otherwise, import all *.py files in it.
|
|
35
|
+
If path is a .py file, import it.
|
|
36
|
+
"""
|
|
37
|
+
p = Path(path)
|
|
38
|
+
if not p.exists():
|
|
39
|
+
return []
|
|
40
|
+
|
|
41
|
+
modules = []
|
|
42
|
+
if p.is_dir():
|
|
43
|
+
if (p / "__init__.py").exists():
|
|
44
|
+
mod = _import_module_from_path(p)
|
|
45
|
+
if mod:
|
|
46
|
+
modules.append(mod)
|
|
47
|
+
else:
|
|
48
|
+
for f in sorted(p.glob("*.py")):
|
|
49
|
+
if f.name != "__init__.py":
|
|
50
|
+
mod = _import_module_from_path(f)
|
|
51
|
+
if mod:
|
|
52
|
+
modules.append(mod)
|
|
53
|
+
elif p.is_file() and p.suffix == ".py":
|
|
54
|
+
mod = _import_module_from_path(p)
|
|
55
|
+
if mod:
|
|
56
|
+
modules.append(mod)
|
|
57
|
+
|
|
58
|
+
return modules
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def _import_module_from_path(file_path: Path) -> Any:
|
|
62
|
+
"""Import a module from a file path dynamically."""
|
|
63
|
+
import sys
|
|
64
|
+
# Use parent directory name if it's an __init__.py file for package names
|
|
65
|
+
if file_path.name == "__init__.py":
|
|
66
|
+
module_name = file_path.parent.name
|
|
67
|
+
elif file_path.is_dir():
|
|
68
|
+
module_name = file_path.name
|
|
69
|
+
file_path = file_path / "__init__.py"
|
|
70
|
+
else:
|
|
71
|
+
module_name = file_path.stem
|
|
72
|
+
|
|
73
|
+
spec = importlib.util.spec_from_file_location(module_name, str(file_path))
|
|
74
|
+
if spec and spec.loader:
|
|
75
|
+
module = importlib.util.module_from_spec(spec)
|
|
76
|
+
sys.modules[module_name] = module
|
|
77
|
+
spec.loader.exec_module(module)
|
|
78
|
+
return module
|
|
79
|
+
|
|
80
|
+
return None
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _build_redis_url(config: RuntimeConfig) -> str:
|
|
84
|
+
"""Build Redis URL from config or environment."""
|
|
85
|
+
if os.getenv("REDIS_URL"):
|
|
86
|
+
return os.getenv("REDIS_URL")
|
|
87
|
+
|
|
88
|
+
if config.backend.dramatiq and config.backend.dramatiq.redis_url:
|
|
89
|
+
return config.backend.dramatiq.redis_url
|
|
90
|
+
|
|
91
|
+
return "redis://localhost:6379/0"
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def initialize_runtime(
|
|
95
|
+
config: RuntimeConfig,
|
|
96
|
+
auto_import: Optional[List[Union[str, Path]]] = None
|
|
97
|
+
) -> Runtime:
|
|
98
|
+
"""
|
|
99
|
+
Initialize the framework runtime: configuration, clients, and task backend.
|
|
100
|
+
|
|
101
|
+
This is the single entry point for setting up the producer side (backend submission)
|
|
102
|
+
and registry initialization.
|
|
103
|
+
"""
|
|
104
|
+
# 1. Init Client Registry
|
|
105
|
+
from beamflow_clients.registry import init_registry
|
|
106
|
+
init_registry(config)
|
|
107
|
+
|
|
108
|
+
# 2. Handle auto-imports from paths
|
|
109
|
+
imported_modules = []
|
|
110
|
+
if auto_import:
|
|
111
|
+
for path in auto_import:
|
|
112
|
+
mods = _auto_import_path(path)
|
|
113
|
+
if mods:
|
|
114
|
+
imported_modules.extend(mods)
|
|
115
|
+
|
|
116
|
+
# 3. Configure task backend based on config.backend.type
|
|
117
|
+
backend: Optional[TaskBackend] = None
|
|
118
|
+
if config.backend.type == BackendType.DRAMATIQ:
|
|
119
|
+
redis_url = _build_redis_url(config)
|
|
120
|
+
broker = RedisBroker(url=redis_url)
|
|
121
|
+
dramatiq.set_broker(broker)
|
|
122
|
+
|
|
123
|
+
# Dramatiq's built-in middlewares (like Prometheus) expect the process_boot
|
|
124
|
+
# signal to be emitted to initialize their internal states (like prometheus Counters).
|
|
125
|
+
# We must emit this manually here because we are initializing the broker programmatically
|
|
126
|
+
# (not via dramatiq CLI), ensuring producers (API) also initialize metrics.
|
|
127
|
+
broker.emit_after("process_boot")
|
|
128
|
+
|
|
129
|
+
# Set beamflow_lib.queue.backend
|
|
130
|
+
backend = set_backend(DramatiqBackend())
|
|
131
|
+
elif config.backend.type == BackendType.ASYNC:
|
|
132
|
+
from beamflow_lib.queue.asyncio_backend import AsyncioBackend
|
|
133
|
+
backend = set_backend(AsyncioBackend())
|
|
134
|
+
elif config.backend.type == BackendType.MANAGED:
|
|
135
|
+
# Managed backend might require environment variables for its URL
|
|
136
|
+
from beamflow_lib.queue.backends.managed_tasks import ManagedTasksBackend
|
|
137
|
+
backend = set_backend(ManagedTasksBackend(
|
|
138
|
+
api_url=os.getenv("MANAGED_API_URL", ""),
|
|
139
|
+
auth_token=os.getenv("MANAGED_AUTH_TOKEN", ""),
|
|
140
|
+
service_url=os.getenv("MANAGED_SERVICE_URL"),
|
|
141
|
+
))
|
|
142
|
+
else:
|
|
143
|
+
raise ValueError(f"Unsupported backend: {config.backend.type}")
|
|
144
|
+
|
|
145
|
+
return Runtime(backend=backend, config=config, imported_modules=imported_modules)
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def build_worker_runtime(
|
|
150
|
+
config: RuntimeConfig,
|
|
151
|
+
auto_import: Optional[List[Union[str, Path]]] = None
|
|
152
|
+
) -> Runtime:
|
|
153
|
+
"""
|
|
154
|
+
Configure runtime for worker service.
|
|
155
|
+
|
|
156
|
+
Deprecated: Use initialize_runtime(config) and then get a consumer.
|
|
157
|
+
"""
|
|
158
|
+
# Import task modules is now out of scope for RuntimeConfig,
|
|
159
|
+
# and they should be passed as auto_import paths or imported manually.
|
|
160
|
+
|
|
161
|
+
return initialize_runtime(config, auto_import=auto_import)
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Worker bootstrap utilities.
|
|
3
|
+
|
|
4
|
+
Provides functions to configure the Dramatiq broker with framework middleware.
|
|
5
|
+
"""
|
|
6
|
+
import dramatiq
|
|
7
|
+
from dramatiq.brokers.redis import RedisBroker
|
|
8
|
+
from beamflow_lib.config.runtime_config import RuntimeConfig, BackendType
|
|
9
|
+
from beamflow_lib.queue.backend import set_backend
|
|
10
|
+
from .dramatiq_backend import DramatiqBackend, FrameworkContextMiddleware
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def _build_redis_url(config: RuntimeConfig) -> str:
|
|
14
|
+
"""Build Redis URL from config or environment."""
|
|
15
|
+
if os.getenv("REDIS_URL"):
|
|
16
|
+
return os.getenv("REDIS_URL")
|
|
17
|
+
|
|
18
|
+
if config.backend.dramatiq and config.backend.dramatiq.redis_url:
|
|
19
|
+
return config.backend.dramatiq.redis_url
|
|
20
|
+
|
|
21
|
+
return "redis://localhost:6379/0"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def configure_worker(config: RuntimeConfig):
|
|
25
|
+
"""
|
|
26
|
+
Configure Dramatiq broker with context middleware.
|
|
27
|
+
|
|
28
|
+
This sets up the broker and middleware for worker processes.
|
|
29
|
+
Note: This is a lower-level function. For most cases, use
|
|
30
|
+
`configure_worker_process()` from entrypoint.py instead.
|
|
31
|
+
|
|
32
|
+
Args:
|
|
33
|
+
config: RuntimeConfig with worker and redis configuration
|
|
34
|
+
|
|
35
|
+
Returns:
|
|
36
|
+
The configured broker instance
|
|
37
|
+
|
|
38
|
+
Raises:
|
|
39
|
+
ValueError: If backend type is not DRAMATIQ
|
|
40
|
+
"""
|
|
41
|
+
if config.backend.type != BackendType.DRAMATIQ:
|
|
42
|
+
raise ValueError(f"Unsupported backend for worker: {config.backend.type}")
|
|
43
|
+
|
|
44
|
+
redis_url = _build_redis_url(config)
|
|
45
|
+
broker = RedisBroker(url=redis_url)
|
|
46
|
+
broker.add_middleware(FrameworkContextMiddleware())
|
|
47
|
+
dramatiq.set_broker(broker)
|
|
48
|
+
set_backend(DramatiqBackend())
|
|
49
|
+
|
|
50
|
+
# Dramatiq's built-in middlewares (like Prometheus) expect the process_boot
|
|
51
|
+
# signal to be emitted to initialize their internal states (like prometheus Counters).
|
|
52
|
+
# Since we initialize the Worker programmatically (not via dramatiq CLI),
|
|
53
|
+
# we must emit this manually. We emit it here so that producers (API) also
|
|
54
|
+
# initialize the metrics, preventing crashes during enqueue.
|
|
55
|
+
broker.emit_after("process_boot")
|
|
56
|
+
|
|
57
|
+
return broker
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
import threading
|
|
2
|
+
import asyncio
|
|
3
|
+
import logging
|
|
4
|
+
import dramatiq
|
|
5
|
+
from beamflow_lib.pipelines.consumer import ConsumerRunner
|
|
6
|
+
|
|
7
|
+
logger = logging.getLogger(__name__)
|
|
8
|
+
|
|
9
|
+
class ConsumerMiddleware(dramatiq.Middleware):
|
|
10
|
+
def __init__(self):
|
|
11
|
+
self._loop: asyncio.AbstractEventLoop = None
|
|
12
|
+
self._thread: threading.Thread = None
|
|
13
|
+
self._runner: ConsumerRunner = None
|
|
14
|
+
self._stop_event = threading.Event()
|
|
15
|
+
|
|
16
|
+
def after_worker_boot(self, broker, worker):
|
|
17
|
+
"""Called when the worker process starts."""
|
|
18
|
+
logger.info("Starting ConsumerRunner thread...")
|
|
19
|
+
|
|
20
|
+
self._thread = threading.Thread(target=self._run_loop, daemon=True, name="ConsumerRunnerThread")
|
|
21
|
+
self._thread.start()
|
|
22
|
+
|
|
23
|
+
def before_worker_shutdown(self, broker, worker):
|
|
24
|
+
"""Called when the worker process is shutting down."""
|
|
25
|
+
logger.info("Stopping ConsumerRunner thread...")
|
|
26
|
+
if self._loop:
|
|
27
|
+
asyncio.run_coroutine_threadsafe(self._runner.stop(), self._loop)
|
|
28
|
+
# Give it a moment to stop
|
|
29
|
+
self._thread.join(timeout=5.0)
|
|
30
|
+
|
|
31
|
+
def _run_loop(self):
|
|
32
|
+
"""Runs the asyncio loop in a separate thread."""
|
|
33
|
+
logger.info("ConsumerRunner loop started.")
|
|
34
|
+
self._loop = asyncio.new_event_loop()
|
|
35
|
+
asyncio.set_event_loop(self._loop)
|
|
36
|
+
|
|
37
|
+
self._runner = ConsumerRunner()
|
|
38
|
+
|
|
39
|
+
self._loop.run_until_complete(self._runner.start())
|
|
40
|
+
|
|
41
|
+
# Keep loop running until stopped
|
|
42
|
+
try:
|
|
43
|
+
self._loop.run_forever()
|
|
44
|
+
except Exception as e:
|
|
45
|
+
logger.error(f"ConsumerRunner loop error: {e}")
|
|
46
|
+
finally:
|
|
47
|
+
self._loop.close()
|
|
48
|
+
logger.info("ConsumerRunner loop closed.")
|