mirror of
https://github.com/frappe/gunicorn.git
synced 2026-07-07 13:11:30 +08:00
feat(dirty): add per-app worker allocation for memory optimization
Allow dirty apps to specify how many workers should load them, enabling significant memory savings for heavy applications like ML models. - Add `workers` class attribute to DirtyApp (None = all workers) - Add `parse_dirty_app_spec()` to parse "module:Class:N" format - Add `DirtyNoWorkersAvailableError` for app-specific error handling - Update DirtyArbiter with per-app worker tracking and routing - Maintain backward compatibility when no dirty_apps configured Example: 8 workers x 10GB model = 80GB RAM needed With workers=2: 2 x 10GB = 20GB RAM (75% savings) Configuration formats: - Class attribute: `workers = 2` on DirtyApp subclass - Config format: `module:class:N` (e.g., `myapp.ml:HugeModel:2`)
This commit is contained in:
parent
1591a6c773
commit
8559854b4f
@ -2885,23 +2885,47 @@ class DirtyApps(Setting):
|
|||||||
desc = """\
|
desc = """\
|
||||||
Dirty applications to load in the dirty worker pool.
|
Dirty applications to load in the dirty worker pool.
|
||||||
|
|
||||||
A list of application paths in pattern ``$(MODULE_NAME):$(CLASS_NAME)``.
|
A list of application paths in one of these formats:
|
||||||
|
|
||||||
|
- ``$(MODULE_NAME):$(CLASS_NAME)`` - all workers load this app
|
||||||
|
- ``$(MODULE_NAME):$(CLASS_NAME):$(N)`` - only N workers load this app
|
||||||
|
|
||||||
Each dirty app must be a class that inherits from ``DirtyApp`` base class
|
Each dirty app must be a class that inherits from ``DirtyApp`` base class
|
||||||
and implements the ``init()``, ``__call__()``, and ``close()`` methods.
|
and implements the ``init()``, ``__call__()``, and ``close()`` methods.
|
||||||
|
|
||||||
Example::
|
Example::
|
||||||
|
|
||||||
dirty_apps = [
|
dirty_apps = [
|
||||||
"myapp.ml:MLApp",
|
"myapp.ml:MLApp", # All workers load this
|
||||||
"myapp.images:ImageApp",
|
"myapp.images:ImageApp", # All workers load this
|
||||||
|
"myapp.heavy:HugeModel:2", # Only 2 workers load this
|
||||||
]
|
]
|
||||||
|
|
||||||
|
The per-app worker limit is useful for memory-intensive applications
|
||||||
|
like large ML models. Instead of all 8 workers loading a 10GB model
|
||||||
|
(80GB total), you can limit it to 2 workers (20GB total).
|
||||||
|
|
||||||
|
Alternatively, you can set the ``workers`` class attribute on your
|
||||||
|
DirtyApp subclass::
|
||||||
|
|
||||||
|
class HugeModelApp(DirtyApp):
|
||||||
|
workers = 2 # Only 2 workers load this app
|
||||||
|
|
||||||
|
def init(self):
|
||||||
|
self.model = load_10gb_model()
|
||||||
|
|
||||||
|
Note: The config format (``module:Class:N``) takes precedence over
|
||||||
|
the class attribute if both are specified.
|
||||||
|
|
||||||
Dirty apps are loaded once when the dirty worker starts and persist
|
Dirty apps are loaded once when the dirty worker starts and persist
|
||||||
in memory for the lifetime of the worker. This is ideal for loading
|
in memory for the lifetime of the worker. This is ideal for loading
|
||||||
ML models, database connection pools, or other stateful resources
|
ML models, database connection pools, or other stateful resources
|
||||||
that are expensive to initialize.
|
that are expensive to initialize.
|
||||||
|
|
||||||
.. versionadded:: 25.0.0
|
.. versionadded:: 25.0.0
|
||||||
|
|
||||||
|
.. versionchanged:: 25.1.0
|
||||||
|
Added per-app worker allocation via ``:N`` format suffix.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@ -72,12 +72,29 @@ class DirtyApp:
|
|||||||
self.models[name] = load_model(name)
|
self.models[name] = load_model(name)
|
||||||
return {"loaded": True, "name": name}
|
return {"loaded": True, "name": name}
|
||||||
|
|
||||||
|
Worker Allocation
|
||||||
|
-----------------
|
||||||
|
By default, all dirty workers load all apps. For apps that consume
|
||||||
|
significant memory (like large ML models), you can limit how many
|
||||||
|
workers load the app by setting the ``workers`` class attribute::
|
||||||
|
|
||||||
|
class HeavyModelApp(DirtyApp):
|
||||||
|
workers = 2 # Only 2 workers will load this app
|
||||||
|
|
||||||
|
def init(self):
|
||||||
|
self.model = load_10gb_model()
|
||||||
|
|
||||||
Subclasses should implement:
|
Subclasses should implement:
|
||||||
- init(): Called once at worker startup to initialize resources
|
- init(): Called once at worker startup to initialize resources
|
||||||
- __call__(action, *args, **kwargs): Handle requests from HTTP workers
|
- __call__(action, *args, **kwargs): Handle requests from HTTP workers
|
||||||
- close(): Called at worker shutdown to cleanup resources
|
- close(): Called at worker shutdown to cleanup resources
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
# Number of workers that should load this app.
|
||||||
|
# None means all workers (default, backward compatible).
|
||||||
|
# Set to an integer to limit how many workers load this app.
|
||||||
|
workers = None
|
||||||
|
|
||||||
def init(self):
|
def init(self):
|
||||||
"""
|
"""
|
||||||
Initialize the application.
|
Initialize the application.
|
||||||
@ -120,6 +137,81 @@ class DirtyApp:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def parse_dirty_app_spec(spec):
|
||||||
|
"""
|
||||||
|
Parse a dirty app specification.
|
||||||
|
|
||||||
|
Supports two formats:
|
||||||
|
- ``"module:Class"`` - standard format, all workers load the app
|
||||||
|
- ``"module:Class:N"`` - worker-limited format, only N workers load the app
|
||||||
|
|
||||||
|
Args:
|
||||||
|
spec: The app specification string
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (import_path, worker_count)
|
||||||
|
- import_path: The "module:Class" part for importing
|
||||||
|
- worker_count: Integer limit or None for all workers
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
DirtyAppError: If the spec format is invalid or worker_count is < 1
|
||||||
|
|
||||||
|
Examples::
|
||||||
|
|
||||||
|
>>> parse_dirty_app_spec("myapp:App")
|
||||||
|
("myapp:App", None)
|
||||||
|
|
||||||
|
>>> parse_dirty_app_spec("myapp:App:2")
|
||||||
|
("myapp:App", 2)
|
||||||
|
|
||||||
|
>>> parse_dirty_app_spec("myapp.sub:App:1")
|
||||||
|
("myapp.sub:App", 1)
|
||||||
|
"""
|
||||||
|
if ':' not in spec:
|
||||||
|
raise DirtyAppError(
|
||||||
|
f"Invalid import path format: {spec}. "
|
||||||
|
f"Expected 'module.path:ClassName' or 'module.path:ClassName:N'",
|
||||||
|
app_path=spec
|
||||||
|
)
|
||||||
|
|
||||||
|
parts = spec.split(':')
|
||||||
|
|
||||||
|
# Standard format: "module:Class" or "module.sub:Class"
|
||||||
|
if len(parts) == 2:
|
||||||
|
return (spec, None)
|
||||||
|
|
||||||
|
# Worker-limited format: "module:Class:N"
|
||||||
|
if len(parts) == 3:
|
||||||
|
module_path, class_name, count_str = parts
|
||||||
|
import_path = f"{module_path}:{class_name}"
|
||||||
|
|
||||||
|
# Validate the worker count
|
||||||
|
try:
|
||||||
|
worker_count = int(count_str)
|
||||||
|
except ValueError:
|
||||||
|
raise DirtyAppError(
|
||||||
|
f"Invalid worker count in spec: {spec}. "
|
||||||
|
f"Expected integer, got '{count_str}'",
|
||||||
|
app_path=spec
|
||||||
|
)
|
||||||
|
|
||||||
|
if worker_count < 1:
|
||||||
|
raise DirtyAppError(
|
||||||
|
f"Invalid worker count in spec: {spec}. "
|
||||||
|
f"Worker count must be >= 1, got {worker_count}",
|
||||||
|
app_path=spec
|
||||||
|
)
|
||||||
|
|
||||||
|
return (import_path, worker_count)
|
||||||
|
|
||||||
|
# Too many colons
|
||||||
|
raise DirtyAppError(
|
||||||
|
f"Invalid import path format: {spec}. "
|
||||||
|
f"Expected 'module.path:ClassName' or 'module.path:ClassName:N'",
|
||||||
|
app_path=spec
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def load_dirty_app(import_path):
|
def load_dirty_app(import_path):
|
||||||
"""
|
"""
|
||||||
Load a dirty app class from an import path.
|
Load a dirty app class from an import path.
|
||||||
|
|||||||
@ -19,7 +19,13 @@ import time
|
|||||||
|
|
||||||
from gunicorn import util
|
from gunicorn import util
|
||||||
|
|
||||||
from .errors import DirtyError, DirtyTimeoutError, DirtyWorkerError
|
from .app import parse_dirty_app_spec
|
||||||
|
from .errors import (
|
||||||
|
DirtyError,
|
||||||
|
DirtyNoWorkersAvailableError,
|
||||||
|
DirtyTimeoutError,
|
||||||
|
DirtyWorkerError,
|
||||||
|
)
|
||||||
from .protocol import (
|
from .protocol import (
|
||||||
DirtyProtocol,
|
DirtyProtocol,
|
||||||
make_error_response,
|
make_error_response,
|
||||||
@ -79,6 +85,100 @@ class DirtyArbiter:
|
|||||||
self._loop = None
|
self._loop = None
|
||||||
self._pending_requests = {} # request_id -> Future
|
self._pending_requests = {} # request_id -> Future
|
||||||
|
|
||||||
|
# Per-app worker allocation tracking
|
||||||
|
# Maps import_path -> {import_path, worker_count, original_spec}
|
||||||
|
self.app_specs = {}
|
||||||
|
# Maps import_path -> set of worker PIDs that have loaded the app
|
||||||
|
self.app_worker_map = {}
|
||||||
|
# Maps worker_pid -> list of import_paths loaded by this worker
|
||||||
|
self.worker_app_map = {}
|
||||||
|
# Per-app round-robin indices for routing
|
||||||
|
self._app_rr_indices = {}
|
||||||
|
# Queue of app lists from dead workers to respawn with same apps
|
||||||
|
self._pending_respawns = []
|
||||||
|
|
||||||
|
# Parse app specs on init
|
||||||
|
self._parse_app_specs()
|
||||||
|
|
||||||
|
def _parse_app_specs(self):
|
||||||
|
"""
|
||||||
|
Parse all app specifications from config.
|
||||||
|
|
||||||
|
Populates self.app_specs with parsed information about each app,
|
||||||
|
including the import path and worker count limits.
|
||||||
|
"""
|
||||||
|
for spec in self.cfg.dirty_apps:
|
||||||
|
import_path, worker_count = parse_dirty_app_spec(spec)
|
||||||
|
self.app_specs[import_path] = {
|
||||||
|
'import_path': import_path,
|
||||||
|
'worker_count': worker_count,
|
||||||
|
'original_spec': spec,
|
||||||
|
}
|
||||||
|
# Initialize the app_worker_map for this app
|
||||||
|
self.app_worker_map[import_path] = set()
|
||||||
|
|
||||||
|
def _get_apps_for_new_worker(self):
|
||||||
|
"""
|
||||||
|
Determine which apps a new worker should load.
|
||||||
|
|
||||||
|
Returns a list of import paths for apps that need more workers.
|
||||||
|
Apps with workers=None (all workers) are always included.
|
||||||
|
Apps with worker limits are included only if they haven't
|
||||||
|
reached their limit yet.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of import paths to load, or empty list if no apps need workers
|
||||||
|
"""
|
||||||
|
app_paths = []
|
||||||
|
|
||||||
|
for import_path, spec in self.app_specs.items():
|
||||||
|
worker_count = spec['worker_count']
|
||||||
|
current_workers = len(self.app_worker_map.get(import_path, set()))
|
||||||
|
|
||||||
|
# None means all workers should load this app
|
||||||
|
if worker_count is None:
|
||||||
|
app_paths.append(import_path)
|
||||||
|
# Otherwise check if we've reached the limit
|
||||||
|
elif current_workers < worker_count:
|
||||||
|
app_paths.append(import_path)
|
||||||
|
|
||||||
|
return app_paths
|
||||||
|
|
||||||
|
def _register_worker_apps(self, worker_pid, app_paths):
|
||||||
|
"""
|
||||||
|
Register which apps a worker has loaded.
|
||||||
|
|
||||||
|
Updates both app_worker_map and worker_app_map to track the
|
||||||
|
bidirectional relationship between workers and apps.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
worker_pid: The PID of the worker
|
||||||
|
app_paths: List of app import paths loaded by this worker
|
||||||
|
"""
|
||||||
|
self.worker_app_map[worker_pid] = list(app_paths)
|
||||||
|
|
||||||
|
for app_path in app_paths:
|
||||||
|
if app_path not in self.app_worker_map:
|
||||||
|
self.app_worker_map[app_path] = set()
|
||||||
|
self.app_worker_map[app_path].add(worker_pid)
|
||||||
|
|
||||||
|
def _unregister_worker(self, worker_pid):
|
||||||
|
"""
|
||||||
|
Unregister a worker's apps when it exits.
|
||||||
|
|
||||||
|
Removes the worker from all tracking maps.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
worker_pid: The PID of the worker to unregister
|
||||||
|
"""
|
||||||
|
# Get the apps this worker had
|
||||||
|
app_paths = self.worker_app_map.pop(worker_pid, [])
|
||||||
|
|
||||||
|
# Remove worker from each app's worker set
|
||||||
|
for app_path in app_paths:
|
||||||
|
if app_path in self.app_worker_map:
|
||||||
|
self.app_worker_map[app_path].discard(worker_pid)
|
||||||
|
|
||||||
def run(self):
|
def run(self):
|
||||||
"""Run the dirty arbiter (blocking call)."""
|
"""Run the dirty arbiter (blocking call)."""
|
||||||
self.pid = os.getpid()
|
self.pid = os.getpid()
|
||||||
@ -256,14 +356,20 @@ class DirtyArbiter:
|
|||||||
client_writer: StreamWriter to send responses to client
|
client_writer: StreamWriter to send responses to client
|
||||||
"""
|
"""
|
||||||
request_id = request.get("id", "unknown")
|
request_id = request.get("id", "unknown")
|
||||||
|
app_path = request.get("app_path")
|
||||||
|
|
||||||
# Find an available worker
|
# Find an available worker (filtered by app if specified)
|
||||||
worker_pid = await self._get_available_worker()
|
worker_pid = await self._get_available_worker(app_path)
|
||||||
if worker_pid is None:
|
if worker_pid is None:
|
||||||
response = make_error_response(
|
# Distinguish between no workers at all vs. no workers for this app
|
||||||
request_id,
|
if not self.workers:
|
||||||
DirtyError("No dirty workers available")
|
error = DirtyError("No dirty workers available")
|
||||||
)
|
elif app_path and self.app_specs:
|
||||||
|
# Per-app allocation is configured and no workers have this app
|
||||||
|
error = DirtyNoWorkersAvailableError(app_path)
|
||||||
|
else:
|
||||||
|
error = DirtyError("No dirty workers available")
|
||||||
|
response = make_error_response(request_id, error)
|
||||||
await DirtyProtocol.write_message_async(client_writer, response)
|
await DirtyProtocol.write_message_async(client_writer, response)
|
||||||
return
|
return
|
||||||
|
|
||||||
@ -373,20 +479,47 @@ class DirtyArbiter:
|
|||||||
)
|
)
|
||||||
await DirtyProtocol.write_message_async(client_writer, response)
|
await DirtyProtocol.write_message_async(client_writer, response)
|
||||||
|
|
||||||
async def _get_available_worker(self):
|
async def _get_available_worker(self, app_path=None):
|
||||||
"""
|
"""
|
||||||
Get an available worker PID using round-robin selection.
|
Get an available worker PID using round-robin selection.
|
||||||
|
|
||||||
Distributes requests across all available workers evenly to
|
If app_path is provided, only returns workers that have loaded
|
||||||
maximize throughput when multiple workers are configured.
|
that specific app. Uses per-app round-robin to ensure fair
|
||||||
|
distribution among eligible workers.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
app_path: Optional import path of the target app. If None,
|
||||||
|
returns any worker using global round-robin.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Worker PID or None if no eligible workers are available.
|
||||||
"""
|
"""
|
||||||
worker_pids = list(self.workers.keys())
|
# Determine eligible workers
|
||||||
if not worker_pids:
|
if app_path and self.app_specs:
|
||||||
|
# Per-app allocation is configured - must return a worker
|
||||||
|
# that has this specific app
|
||||||
|
if app_path in self.app_worker_map:
|
||||||
|
eligible_pids = list(self.app_worker_map[app_path])
|
||||||
|
else:
|
||||||
|
# App not known or no workers have it
|
||||||
|
return None
|
||||||
|
else:
|
||||||
|
# No specific app requested, or no app specs configured
|
||||||
|
# (backward compatible) - any worker will do
|
||||||
|
eligible_pids = list(self.workers.keys())
|
||||||
|
|
||||||
|
if not eligible_pids:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Round-robin selection
|
# Per-app round-robin for fairness
|
||||||
self._worker_rr_index = (self._worker_rr_index + 1) % len(worker_pids)
|
if app_path and self.app_specs:
|
||||||
return worker_pids[self._worker_rr_index]
|
idx = self._app_rr_indices.get(app_path, 0)
|
||||||
|
self._app_rr_indices[app_path] = (idx + 1) % len(eligible_pids)
|
||||||
|
else:
|
||||||
|
idx = self._worker_rr_index
|
||||||
|
self._worker_rr_index = (idx + 1) % len(eligible_pids)
|
||||||
|
|
||||||
|
return eligible_pids[idx % len(eligible_pids)]
|
||||||
|
|
||||||
async def _get_worker_connection(self, worker_pid):
|
async def _get_worker_connection(self, worker_pid):
|
||||||
"""Get or create connection to a worker."""
|
"""Get or create connection to a worker."""
|
||||||
@ -424,7 +557,10 @@ class DirtyArbiter:
|
|||||||
|
|
||||||
# Spawn workers if needed
|
# Spawn workers if needed
|
||||||
while self.alive and len(self.workers) < num_workers:
|
while self.alive and len(self.workers) < num_workers:
|
||||||
self.spawn_worker()
|
result = self.spawn_worker()
|
||||||
|
if result is None:
|
||||||
|
# No apps need more workers - stop spawning
|
||||||
|
break
|
||||||
await asyncio.sleep(0.1)
|
await asyncio.sleep(0.1)
|
||||||
|
|
||||||
# Kill excess workers
|
# Kill excess workers
|
||||||
@ -436,7 +572,27 @@ class DirtyArbiter:
|
|||||||
await asyncio.sleep(0.1)
|
await asyncio.sleep(0.1)
|
||||||
|
|
||||||
def spawn_worker(self):
|
def spawn_worker(self):
|
||||||
"""Spawn a new dirty worker."""
|
"""
|
||||||
|
Spawn a new dirty worker.
|
||||||
|
|
||||||
|
Worker app assignment follows these priorities:
|
||||||
|
1. If there are pending respawns (from dead workers), use those apps
|
||||||
|
2. Otherwise, determine apps for a new worker based on allocation
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Worker PID in parent process, or None if no apps need workers
|
||||||
|
"""
|
||||||
|
# Priority 1: Respawn dead worker with same apps
|
||||||
|
if self._pending_respawns:
|
||||||
|
app_paths = self._pending_respawns.pop(0)
|
||||||
|
else:
|
||||||
|
# Priority 2: New worker for initial pool
|
||||||
|
app_paths = self._get_apps_for_new_worker()
|
||||||
|
|
||||||
|
if not app_paths:
|
||||||
|
self.log.warning("No apps need more workers, skipping spawn")
|
||||||
|
return None
|
||||||
|
|
||||||
self.worker_age += 1
|
self.worker_age += 1
|
||||||
socket_path = os.path.join(
|
socket_path = os.path.join(
|
||||||
self.tmpdir, f"worker-{self.worker_age}.sock"
|
self.tmpdir, f"worker-{self.worker_age}.sock"
|
||||||
@ -445,7 +601,7 @@ class DirtyArbiter:
|
|||||||
worker = DirtyWorker(
|
worker = DirtyWorker(
|
||||||
age=self.worker_age,
|
age=self.worker_age,
|
||||||
ppid=self.pid,
|
ppid=self.pid,
|
||||||
app_paths=self.cfg.dirty_apps,
|
app_paths=app_paths, # Only assigned apps, not all apps
|
||||||
cfg=self.cfg,
|
cfg=self.cfg,
|
||||||
log=self.log,
|
log=self.log,
|
||||||
socket_path=socket_path
|
socket_path=socket_path
|
||||||
@ -457,8 +613,13 @@ class DirtyArbiter:
|
|||||||
worker.pid = pid
|
worker.pid = pid
|
||||||
self.workers[pid] = worker
|
self.workers[pid] = worker
|
||||||
self.worker_sockets[pid] = socket_path
|
self.worker_sockets[pid] = socket_path
|
||||||
|
|
||||||
|
# Register which apps this worker has
|
||||||
|
self._register_worker_apps(pid, app_paths)
|
||||||
|
|
||||||
self.cfg.dirty_post_fork(self, worker)
|
self.cfg.dirty_post_fork(self, worker)
|
||||||
self.log.info("Spawned dirty worker (pid: %s)", pid)
|
self.log.info("Spawned dirty worker (pid: %s) with apps: %s",
|
||||||
|
pid, app_paths)
|
||||||
return pid
|
return pid
|
||||||
|
|
||||||
# Child process
|
# Child process
|
||||||
@ -484,7 +645,12 @@ class DirtyArbiter:
|
|||||||
self._cleanup_worker(pid)
|
self._cleanup_worker(pid)
|
||||||
|
|
||||||
def _cleanup_worker(self, pid):
|
def _cleanup_worker(self, pid):
|
||||||
"""Clean up after a worker exits."""
|
"""
|
||||||
|
Clean up after a worker exits.
|
||||||
|
|
||||||
|
Saves the dead worker's app list to pending respawns so the
|
||||||
|
replacement worker gets the same apps.
|
||||||
|
"""
|
||||||
self._close_worker_connection(pid)
|
self._close_worker_connection(pid)
|
||||||
|
|
||||||
# Cancel consumer task
|
# Cancel consumer task
|
||||||
@ -495,6 +661,15 @@ class DirtyArbiter:
|
|||||||
# Remove queue
|
# Remove queue
|
||||||
self.worker_queues.pop(pid, None)
|
self.worker_queues.pop(pid, None)
|
||||||
|
|
||||||
|
# Save dead worker's apps for respawn BEFORE unregistering
|
||||||
|
if pid in self.worker_app_map:
|
||||||
|
dead_apps = list(self.worker_app_map[pid])
|
||||||
|
if dead_apps:
|
||||||
|
self._pending_respawns.append(dead_apps)
|
||||||
|
|
||||||
|
# Now safe to unregister the worker's apps
|
||||||
|
self._unregister_worker(pid)
|
||||||
|
|
||||||
worker = self.workers.pop(pid, None)
|
worker = self.workers.pop(pid, None)
|
||||||
if worker:
|
if worker:
|
||||||
self.cfg.dirty_worker_exit(self, worker)
|
self.cfg.dirty_worker_exit(self, worker)
|
||||||
|
|||||||
@ -46,6 +46,7 @@ class DirtyError(Exception):
|
|||||||
"DirtyWorkerError": DirtyWorkerError,
|
"DirtyWorkerError": DirtyWorkerError,
|
||||||
"DirtyAppError": DirtyAppError,
|
"DirtyAppError": DirtyAppError,
|
||||||
"DirtyAppNotFoundError": DirtyAppNotFoundError,
|
"DirtyAppNotFoundError": DirtyAppNotFoundError,
|
||||||
|
"DirtyNoWorkersAvailableError": DirtyNoWorkersAvailableError,
|
||||||
"DirtyProtocolError": DirtyProtocolError,
|
"DirtyProtocolError": DirtyProtocolError,
|
||||||
}
|
}
|
||||||
error_type = data.get("error_type", "DirtyError")
|
error_type = data.get("error_type", "DirtyError")
|
||||||
@ -70,6 +71,8 @@ class DirtyError(Exception):
|
|||||||
error.app_path = error.details.get("app_path")
|
error.app_path = error.details.get("app_path")
|
||||||
error.action = error.details.get("action")
|
error.action = error.details.get("action")
|
||||||
error.traceback = error.details.get("traceback")
|
error.traceback = error.details.get("traceback")
|
||||||
|
elif error_class == DirtyNoWorkersAvailableError:
|
||||||
|
error.app_path = error.details.get("app_path")
|
||||||
|
|
||||||
return error
|
return error
|
||||||
|
|
||||||
@ -130,6 +133,40 @@ class DirtyAppNotFoundError(DirtyAppError):
|
|||||||
super().__init__(f"Dirty app not found: {app_path}", app_path=app_path)
|
super().__init__(f"Dirty app not found: {app_path}", app_path=app_path)
|
||||||
|
|
||||||
|
|
||||||
|
class DirtyNoWorkersAvailableError(DirtyError):
|
||||||
|
"""
|
||||||
|
Raised when no workers are available for the requested app.
|
||||||
|
|
||||||
|
This exception is raised when a request targets an app that has
|
||||||
|
worker limits configured, and no workers with that app are currently
|
||||||
|
available (e.g., all workers for that app crashed and haven't been
|
||||||
|
respawned yet).
|
||||||
|
|
||||||
|
Web applications can catch this exception to provide graceful
|
||||||
|
degradation, such as queuing requests for retry or showing a
|
||||||
|
maintenance page.
|
||||||
|
|
||||||
|
Example::
|
||||||
|
|
||||||
|
from gunicorn.dirty import get_dirty_client
|
||||||
|
from gunicorn.dirty.errors import DirtyNoWorkersAvailableError
|
||||||
|
|
||||||
|
def my_view(request):
|
||||||
|
client = get_dirty_client()
|
||||||
|
try:
|
||||||
|
result = client.execute("myapp.ml:HeavyModel", "predict", data)
|
||||||
|
except DirtyNoWorkersAvailableError as e:
|
||||||
|
return {"error": "Service temporarily unavailable",
|
||||||
|
"app": e.app_path}
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, app_path, message=None):
|
||||||
|
if message is None:
|
||||||
|
message = f"No workers available for app: {app_path}"
|
||||||
|
super().__init__(message, details={"app_path": app_path})
|
||||||
|
self.app_path = app_path
|
||||||
|
|
||||||
|
|
||||||
class DirtyProtocolError(DirtyError):
|
class DirtyProtocolError(DirtyError):
|
||||||
"""Raised when there is a protocol-level error."""
|
"""Raised when there is a protocol-level error."""
|
||||||
|
|
||||||
|
|||||||
273
tests/dirty/test_per_app_worker_allocation.py
Normal file
273
tests/dirty/test_per_app_worker_allocation.py
Normal file
@ -0,0 +1,273 @@
|
|||||||
|
#
|
||||||
|
# This file is part of gunicorn released under the MIT license.
|
||||||
|
# See the NOTICE for more information.
|
||||||
|
|
||||||
|
"""Integration tests for per-app worker allocation."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from gunicorn.config import Config
|
||||||
|
from gunicorn.dirty.arbiter import DirtyArbiter
|
||||||
|
|
||||||
|
|
||||||
|
class MockLog:
|
||||||
|
"""Mock logger for testing."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.messages = []
|
||||||
|
|
||||||
|
def debug(self, msg, *args):
|
||||||
|
self.messages.append(("debug", msg % args if args else msg))
|
||||||
|
|
||||||
|
def info(self, msg, *args):
|
||||||
|
self.messages.append(("info", msg % args if args else msg))
|
||||||
|
|
||||||
|
def warning(self, msg, *args):
|
||||||
|
self.messages.append(("warning", msg % args if args else msg))
|
||||||
|
|
||||||
|
def error(self, msg, *args):
|
||||||
|
self.messages.append(("error", msg % args if args else msg))
|
||||||
|
|
||||||
|
def critical(self, msg, *args):
|
||||||
|
self.messages.append(("critical", msg % args if args else msg))
|
||||||
|
|
||||||
|
def exception(self, msg, *args):
|
||||||
|
self.messages.append(("exception", msg % args if args else msg))
|
||||||
|
|
||||||
|
def close_on_exec(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def reopen_files(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class TestPerAppWorkerAllocation:
|
||||||
|
"""Integration tests for per-app worker allocation."""
|
||||||
|
|
||||||
|
def test_heavy_app_loaded_on_limited_workers(self):
|
||||||
|
"""App with workers=2 only loaded on 2 of 4 workers."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_workers", 4)
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp", # unlimited
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:2", # limited to 2
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# Simulate spawning 4 workers
|
||||||
|
for i in range(4):
|
||||||
|
apps = arbiter._get_apps_for_new_worker()
|
||||||
|
arbiter._register_worker_apps(1000 + i, apps)
|
||||||
|
|
||||||
|
# Check distribution
|
||||||
|
unlimited_app = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
limited_app = "tests.support_dirty_app:SlowDirtyApp"
|
||||||
|
|
||||||
|
# Unlimited app should be on all 4 workers
|
||||||
|
assert len(arbiter.app_worker_map[unlimited_app]) == 4
|
||||||
|
|
||||||
|
# Limited app should only be on 2 workers
|
||||||
|
assert len(arbiter.app_worker_map[limited_app]) == 2
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_light_app_loaded_on_all_workers(self):
|
||||||
|
"""App with workers=None loaded on all workers."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_workers", 4)
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# Simulate spawning 4 workers
|
||||||
|
for i in range(4):
|
||||||
|
apps = arbiter._get_apps_for_new_worker()
|
||||||
|
arbiter._register_worker_apps(1000 + i, apps)
|
||||||
|
|
||||||
|
# App should be on all 4 workers
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
assert len(arbiter.app_worker_map[app_path]) == 4
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_mixed_apps_correct_distribution(self):
|
||||||
|
"""Mix of limited and unlimited apps distributed correctly."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_workers", 4)
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp", # unlimited
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:1", # limited to 1
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# Simulate spawning 4 workers
|
||||||
|
for i in range(4):
|
||||||
|
apps = arbiter._get_apps_for_new_worker()
|
||||||
|
arbiter._register_worker_apps(1000 + i, apps)
|
||||||
|
|
||||||
|
unlimited_app = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
limited_app = "tests.support_dirty_app:SlowDirtyApp"
|
||||||
|
|
||||||
|
# Unlimited app on all workers
|
||||||
|
assert len(arbiter.app_worker_map[unlimited_app]) == 4
|
||||||
|
|
||||||
|
# Limited app on only 1 worker
|
||||||
|
assert len(arbiter.app_worker_map[limited_app]) == 1
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_request_routing_respects_allocation(self):
|
||||||
|
"""Requests only routed to workers with the target app."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:1",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# Set up workers
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.workers[1002] = "worker2"
|
||||||
|
|
||||||
|
# Worker 1001 has both apps, worker 1002 has only TestDirtyApp
|
||||||
|
arbiter._register_worker_apps(1001, [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp",
|
||||||
|
])
|
||||||
|
arbiter._register_worker_apps(1002, [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
])
|
||||||
|
|
||||||
|
# Request for SlowDirtyApp should only go to worker 1001
|
||||||
|
worker = await arbiter._get_available_worker("tests.support_dirty_app:SlowDirtyApp")
|
||||||
|
assert worker == 1001
|
||||||
|
|
||||||
|
# Request for TestDirtyApp should go to either
|
||||||
|
worker = await arbiter._get_available_worker("tests.support_dirty_app:TestDirtyApp")
|
||||||
|
assert worker in [1001, 1002]
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_worker_crash_app_reassigned_to_new_worker(self):
|
||||||
|
"""When worker dies, new worker gets the app it had."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:1",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
|
||||||
|
# Set up initial workers
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.worker_sockets[1001] = "/tmp/fake1.sock"
|
||||||
|
|
||||||
|
# Worker 1001 has both apps
|
||||||
|
arbiter._register_worker_apps(1001, [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp",
|
||||||
|
])
|
||||||
|
|
||||||
|
# Simulate worker crash
|
||||||
|
arbiter._cleanup_worker(1001)
|
||||||
|
|
||||||
|
# Apps should be queued for respawn
|
||||||
|
assert len(arbiter._pending_respawns) == 1
|
||||||
|
pending_apps = arbiter._pending_respawns[0]
|
||||||
|
assert "tests.support_dirty_app:TestDirtyApp" in pending_apps
|
||||||
|
assert "tests.support_dirty_app:SlowDirtyApp" in pending_apps
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_worker_crash_other_workers_still_serve_app(self):
|
||||||
|
"""When one of two workers dies, other still serves requests."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
|
||||||
|
# Set up two workers for the same app
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.worker_sockets[1001] = "/tmp/fake1.sock"
|
||||||
|
arbiter.workers[1002] = "worker2"
|
||||||
|
arbiter.worker_sockets[1002] = "/tmp/fake2.sock"
|
||||||
|
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
arbiter._register_worker_apps(1002, [app_path])
|
||||||
|
|
||||||
|
# Both workers serve the app
|
||||||
|
assert len(arbiter.app_worker_map[app_path]) == 2
|
||||||
|
|
||||||
|
# Worker 1001 crashes
|
||||||
|
arbiter._cleanup_worker(1001)
|
||||||
|
|
||||||
|
# Worker 1002 still serves requests
|
||||||
|
assert len(arbiter.app_worker_map[app_path]) == 1
|
||||||
|
assert 1002 in arbiter.app_worker_map[app_path]
|
||||||
|
|
||||||
|
worker = await arbiter._get_available_worker(app_path)
|
||||||
|
assert worker == 1002
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_worker_crash_sole_worker_app_unavailable_until_respawn(self):
|
||||||
|
"""When sole worker for app dies, requests fail until respawn."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:1",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
|
||||||
|
# Only one worker for this app
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.worker_sockets[1001] = "/tmp/fake1.sock"
|
||||||
|
|
||||||
|
app_path = "tests.support_dirty_app:SlowDirtyApp"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
|
||||||
|
# Worker crashes
|
||||||
|
arbiter._cleanup_worker(1001)
|
||||||
|
|
||||||
|
# No workers available for the app
|
||||||
|
worker = await arbiter._get_available_worker(app_path)
|
||||||
|
assert worker is None
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_config_format_module_class_n(self):
|
||||||
|
"""Config 'mod:Class:2' correctly limits to 2 workers."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp:2",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# Check parsed spec
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
assert arbiter.app_specs[app_path]["worker_count"] == 2
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
@ -6,7 +6,12 @@
|
|||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from gunicorn.dirty.app import DirtyApp, load_dirty_app, load_dirty_apps
|
from gunicorn.dirty.app import (
|
||||||
|
DirtyApp,
|
||||||
|
load_dirty_app,
|
||||||
|
load_dirty_apps,
|
||||||
|
parse_dirty_app_spec,
|
||||||
|
)
|
||||||
from gunicorn.dirty.errors import DirtyAppError, DirtyAppNotFoundError
|
from gunicorn.dirty.errors import DirtyAppError, DirtyAppNotFoundError
|
||||||
|
|
||||||
|
|
||||||
@ -185,3 +190,108 @@ class TestDirtyAppStateful:
|
|||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
app("compute", 1, 2, operation="invalid")
|
app("compute", 1, 2, operation="invalid")
|
||||||
assert "Unknown operation" in str(exc_info.value)
|
assert "Unknown operation" in str(exc_info.value)
|
||||||
|
|
||||||
|
|
||||||
|
class TestDirtyAppWorkersAttribute:
|
||||||
|
"""Tests for DirtyApp workers class attribute."""
|
||||||
|
|
||||||
|
def test_default_workers_is_none(self):
|
||||||
|
"""Base DirtyApp has workers=None (all workers)."""
|
||||||
|
assert DirtyApp.workers is None
|
||||||
|
|
||||||
|
def test_subclass_can_set_workers(self):
|
||||||
|
"""Subclass can override workers=2."""
|
||||||
|
|
||||||
|
class LimitedApp(DirtyApp):
|
||||||
|
workers = 2
|
||||||
|
|
||||||
|
assert LimitedApp.workers == 2
|
||||||
|
|
||||||
|
def test_workers_inherited_by_default(self):
|
||||||
|
"""Subclass without workers attr inherits None."""
|
||||||
|
|
||||||
|
class InheritedApp(DirtyApp):
|
||||||
|
pass
|
||||||
|
|
||||||
|
assert InheritedApp.workers is None
|
||||||
|
|
||||||
|
def test_instance_has_workers_attribute(self):
|
||||||
|
"""Instance should have access to workers attribute."""
|
||||||
|
app = DirtyApp()
|
||||||
|
assert app.workers is None
|
||||||
|
|
||||||
|
class LimitedApp(DirtyApp):
|
||||||
|
workers = 3
|
||||||
|
|
||||||
|
limited = LimitedApp()
|
||||||
|
assert limited.workers == 3
|
||||||
|
|
||||||
|
|
||||||
|
class TestParseDirtyAppSpec:
|
||||||
|
"""Tests for parse_dirty_app_spec function."""
|
||||||
|
|
||||||
|
def test_standard_format(self):
|
||||||
|
"""'mod:Class' returns ('mod:Class', None)."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod:Class")
|
||||||
|
assert import_path == "mod:Class"
|
||||||
|
assert count is None
|
||||||
|
|
||||||
|
def test_standard_format_with_dots(self):
|
||||||
|
"""'mod.sub.pkg:Class' returns ('mod.sub.pkg:Class', None)."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod.sub.pkg:Class")
|
||||||
|
assert import_path == "mod.sub.pkg:Class"
|
||||||
|
assert count is None
|
||||||
|
|
||||||
|
def test_with_worker_count(self):
|
||||||
|
"""'mod:Class:2' returns ('mod:Class', 2)."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod:Class:2")
|
||||||
|
assert import_path == "mod:Class"
|
||||||
|
assert count == 2
|
||||||
|
|
||||||
|
def test_worker_count_one(self):
|
||||||
|
"""'mod:Class:1' returns ('mod:Class', 1)."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod:Class:1")
|
||||||
|
assert import_path == "mod:Class"
|
||||||
|
assert count == 1
|
||||||
|
|
||||||
|
def test_worker_count_large(self):
|
||||||
|
"""'mod:Class:100' returns ('mod:Class', 100)."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod:Class:100")
|
||||||
|
assert import_path == "mod:Class"
|
||||||
|
assert count == 100
|
||||||
|
|
||||||
|
def test_worker_count_zero_raises(self):
|
||||||
|
"""'mod:Class:0' raises DirtyAppError."""
|
||||||
|
with pytest.raises(DirtyAppError) as exc_info:
|
||||||
|
parse_dirty_app_spec("mod:Class:0")
|
||||||
|
assert "must be >= 1" in str(exc_info.value)
|
||||||
|
|
||||||
|
def test_worker_count_negative_raises(self):
|
||||||
|
"""'mod:Class:-1' raises DirtyAppError."""
|
||||||
|
with pytest.raises(DirtyAppError) as exc_info:
|
||||||
|
parse_dirty_app_spec("mod:Class:-1")
|
||||||
|
assert "must be >= 1" in str(exc_info.value)
|
||||||
|
|
||||||
|
def test_non_numeric_raises(self):
|
||||||
|
"""'mod:Class:abc' raises DirtyAppError."""
|
||||||
|
with pytest.raises(DirtyAppError) as exc_info:
|
||||||
|
parse_dirty_app_spec("mod:Class:abc")
|
||||||
|
assert "Expected integer" in str(exc_info.value)
|
||||||
|
|
||||||
|
def test_no_colon_raises(self):
|
||||||
|
"""'mod.Class' (no colon) raises DirtyAppError."""
|
||||||
|
with pytest.raises(DirtyAppError) as exc_info:
|
||||||
|
parse_dirty_app_spec("mod.Class")
|
||||||
|
assert "Invalid import path format" in str(exc_info.value)
|
||||||
|
|
||||||
|
def test_too_many_colons_raises(self):
|
||||||
|
"""'mod:Class:2:extra' raises DirtyAppError."""
|
||||||
|
with pytest.raises(DirtyAppError) as exc_info:
|
||||||
|
parse_dirty_app_spec("mod:Class:2:extra")
|
||||||
|
assert "Invalid import path format" in str(exc_info.value)
|
||||||
|
|
||||||
|
def test_dotted_module_with_count(self):
|
||||||
|
"""'mod.sub:Class:2' handles dots correctly."""
|
||||||
|
import_path, count = parse_dirty_app_spec("mod.sub:Class:2")
|
||||||
|
assert import_path == "mod.sub:Class"
|
||||||
|
assert count == 2
|
||||||
|
|||||||
@ -1144,3 +1144,275 @@ class TestDirtyArbiterQueueBehavior:
|
|||||||
assert task2.done()
|
assert task2.done()
|
||||||
|
|
||||||
arbiter._cleanup_sync()
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
|
||||||
|
class TestDirtyArbiterAppTracking:
|
||||||
|
"""Tests for per-app worker tracking."""
|
||||||
|
|
||||||
|
def test_parse_app_specs_standard_format(self):
|
||||||
|
"""All standard format apps have worker_count=None."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
assert len(arbiter.app_specs) == 2
|
||||||
|
assert arbiter.app_specs["tests.support_dirty_app:TestDirtyApp"]["worker_count"] is None
|
||||||
|
assert arbiter.app_specs["tests.support_dirty_app:SlowDirtyApp"]["worker_count"] is None
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_parse_app_specs_with_worker_count(self):
|
||||||
|
"""Apps with :N have correct worker_count."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:2",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
assert arbiter.app_specs["tests.support_dirty_app:TestDirtyApp"]["worker_count"] is None
|
||||||
|
assert arbiter.app_specs["tests.support_dirty_app:SlowDirtyApp"]["worker_count"] == 2
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_get_apps_for_new_worker_all_standard(self):
|
||||||
|
"""All apps returned when all have workers=None."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp",
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp",
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
apps = arbiter._get_apps_for_new_worker()
|
||||||
|
assert len(apps) == 2
|
||||||
|
assert "tests.support_dirty_app:TestDirtyApp" in apps
|
||||||
|
assert "tests.support_dirty_app:SlowDirtyApp" in apps
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_get_apps_for_new_worker_respects_limit(self):
|
||||||
|
"""App with workers=2 stops assigning after 2 workers."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", [
|
||||||
|
"tests.support_dirty_app:TestDirtyApp", # unlimited
|
||||||
|
"tests.support_dirty_app:SlowDirtyApp:2", # limited to 2
|
||||||
|
])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
# First worker should get both apps
|
||||||
|
apps1 = arbiter._get_apps_for_new_worker()
|
||||||
|
assert len(apps1) == 2
|
||||||
|
arbiter._register_worker_apps(1001, apps1)
|
||||||
|
|
||||||
|
# Second worker should get both apps
|
||||||
|
apps2 = arbiter._get_apps_for_new_worker()
|
||||||
|
assert len(apps2) == 2
|
||||||
|
arbiter._register_worker_apps(1002, apps2)
|
||||||
|
|
||||||
|
# Third worker should only get unlimited app
|
||||||
|
apps3 = arbiter._get_apps_for_new_worker()
|
||||||
|
assert len(apps3) == 1
|
||||||
|
assert "tests.support_dirty_app:TestDirtyApp" in apps3
|
||||||
|
assert "tests.support_dirty_app:SlowDirtyApp" not in apps3
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_register_worker_apps_updates_both_maps(self):
|
||||||
|
"""Both app_worker_map and worker_app_map updated."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
|
||||||
|
# Check app_worker_map
|
||||||
|
assert 1001 in arbiter.app_worker_map[app_path]
|
||||||
|
|
||||||
|
# Check worker_app_map
|
||||||
|
assert app_path in arbiter.worker_app_map[1001]
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_unregister_worker_cleans_both_maps(self):
|
||||||
|
"""Worker removal updates both maps correctly."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
|
||||||
|
# Verify registered
|
||||||
|
assert 1001 in arbiter.app_worker_map[app_path]
|
||||||
|
assert 1001 in arbiter.worker_app_map
|
||||||
|
|
||||||
|
# Unregister
|
||||||
|
arbiter._unregister_worker(1001)
|
||||||
|
|
||||||
|
# Verify cleaned up
|
||||||
|
assert 1001 not in arbiter.app_worker_map[app_path]
|
||||||
|
assert 1001 not in arbiter.worker_app_map
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
|
||||||
|
class TestDirtyArbiterSpawnWorkerPerApp:
|
||||||
|
"""Tests for spawn_worker with per-app allocation."""
|
||||||
|
|
||||||
|
def test_cleanup_worker_queues_apps_for_respawn(self):
|
||||||
|
"""Dead worker's apps added to _pending_respawns."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
|
||||||
|
# Simulate worker registration
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter.workers[1001] = "fake_worker"
|
||||||
|
arbiter.worker_sockets[1001] = "/tmp/fake.sock"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
|
||||||
|
# Cleanup should queue apps for respawn
|
||||||
|
assert len(arbiter._pending_respawns) == 0
|
||||||
|
arbiter._cleanup_worker(1001)
|
||||||
|
assert len(arbiter._pending_respawns) == 1
|
||||||
|
assert app_path in arbiter._pending_respawns[0]
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
def test_pending_respawns_cleared_after_spawn(self):
|
||||||
|
"""Pending respawns consumed when spawning new worker."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
|
||||||
|
# Add pending respawn
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter._pending_respawns.append([app_path])
|
||||||
|
|
||||||
|
# Get apps for new worker should use pending first
|
||||||
|
# But since spawn_worker forks, we test the logic directly
|
||||||
|
assert len(arbiter._pending_respawns) == 1
|
||||||
|
|
||||||
|
# When spawn_worker pops from pending_respawns
|
||||||
|
apps = arbiter._pending_respawns.pop(0)
|
||||||
|
assert apps == [app_path]
|
||||||
|
assert len(arbiter._pending_respawns) == 0
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
|
||||||
|
class TestDirtyArbiterRoutingPerApp:
|
||||||
|
"""Tests for app-aware routing."""
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_available_worker_no_filter(self):
|
||||||
|
"""Without app_path, returns any worker round-robin."""
|
||||||
|
cfg = Config()
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.workers[1002] = "worker2"
|
||||||
|
|
||||||
|
# Should return workers in round-robin
|
||||||
|
w1 = await arbiter._get_available_worker()
|
||||||
|
w2 = await arbiter._get_available_worker()
|
||||||
|
|
||||||
|
assert w1 in [1001, 1002]
|
||||||
|
assert w2 in [1001, 1002]
|
||||||
|
# They should be different (round-robin)
|
||||||
|
if len(arbiter.workers) >= 2:
|
||||||
|
assert w1 != w2 or len(arbiter.workers) == 1
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_available_worker_with_app_filter(self):
|
||||||
|
"""With app_path, returns only workers that have it."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
arbiter.workers[1002] = "worker2"
|
||||||
|
|
||||||
|
# Only register 1001 for the app
|
||||||
|
app_path = "tests.support_dirty_app:TestDirtyApp"
|
||||||
|
arbiter._register_worker_apps(1001, [app_path])
|
||||||
|
|
||||||
|
# Should only return 1001
|
||||||
|
worker = await arbiter._get_available_worker(app_path)
|
||||||
|
assert worker == 1001
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_available_worker_app_no_workers_returns_none(self):
|
||||||
|
"""Returns None if no workers have the app."""
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
|
||||||
|
# Worker 1001 has no apps registered - request for unknown app returns None
|
||||||
|
worker = await arbiter._get_available_worker("unknown:App")
|
||||||
|
assert worker is None
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_route_request_app_not_loaded_error(self):
|
||||||
|
"""Error response when no worker has the app."""
|
||||||
|
from gunicorn.dirty.protocol import DirtyProtocol
|
||||||
|
|
||||||
|
cfg = Config()
|
||||||
|
cfg.set("dirty_apps", ["tests.support_dirty_app:TestDirtyApp"])
|
||||||
|
log = MockLog()
|
||||||
|
|
||||||
|
arbiter = DirtyArbiter(cfg=cfg, log=log)
|
||||||
|
arbiter.pid = 12345
|
||||||
|
arbiter.workers[1001] = "worker1"
|
||||||
|
# No apps registered for this worker (worker exists but has no apps)
|
||||||
|
|
||||||
|
request = make_request(
|
||||||
|
request_id="test-123",
|
||||||
|
app_path="unknown:App",
|
||||||
|
action="test"
|
||||||
|
)
|
||||||
|
|
||||||
|
writer = MockStreamWriter()
|
||||||
|
await arbiter.route_request(request, writer)
|
||||||
|
|
||||||
|
assert len(writer.messages) == 1
|
||||||
|
response = writer.messages[0]
|
||||||
|
assert response["type"] == DirtyProtocol.MSG_TYPE_ERROR
|
||||||
|
assert "No workers available for app" in response["error"]["message"]
|
||||||
|
assert response["error"]["error_type"] == "DirtyNoWorkersAvailableError"
|
||||||
|
|
||||||
|
arbiter._cleanup_sync()
|
||||||
|
|||||||
76
tests/test_dirty_errors.py
Normal file
76
tests/test_dirty_errors.py
Normal file
@ -0,0 +1,76 @@
|
|||||||
|
#
|
||||||
|
# This file is part of gunicorn released under the MIT license.
|
||||||
|
# See the NOTICE for more information.
|
||||||
|
|
||||||
|
"""Tests for dirty errors module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from gunicorn.dirty.errors import (
|
||||||
|
DirtyError,
|
||||||
|
DirtyNoWorkersAvailableError,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestDirtyNoWorkersAvailableError:
|
||||||
|
"""Tests for DirtyNoWorkersAvailableError exception."""
|
||||||
|
|
||||||
|
def test_error_contains_app_path(self):
|
||||||
|
"""Error includes the app_path."""
|
||||||
|
error = DirtyNoWorkersAvailableError("myapp:Model")
|
||||||
|
assert error.app_path == "myapp:Model"
|
||||||
|
assert "myapp:Model" in str(error)
|
||||||
|
assert "No workers available" in str(error)
|
||||||
|
|
||||||
|
def test_error_with_custom_message(self):
|
||||||
|
"""Error can have a custom message."""
|
||||||
|
error = DirtyNoWorkersAvailableError(
|
||||||
|
"myapp:Model",
|
||||||
|
message="Custom: no workers for heavy model"
|
||||||
|
)
|
||||||
|
assert error.app_path == "myapp:Model"
|
||||||
|
assert "Custom: no workers" in str(error)
|
||||||
|
|
||||||
|
def test_error_serialization_roundtrip(self):
|
||||||
|
"""Error survives to_dict/from_dict cycle."""
|
||||||
|
original = DirtyNoWorkersAvailableError("myapp.ml:HugeModel")
|
||||||
|
|
||||||
|
# Serialize
|
||||||
|
data = original.to_dict()
|
||||||
|
assert data["error_type"] == "DirtyNoWorkersAvailableError"
|
||||||
|
assert data["details"]["app_path"] == "myapp.ml:HugeModel"
|
||||||
|
|
||||||
|
# Deserialize
|
||||||
|
restored = DirtyError.from_dict(data)
|
||||||
|
assert isinstance(restored, DirtyNoWorkersAvailableError)
|
||||||
|
assert restored.app_path == "myapp.ml:HugeModel"
|
||||||
|
assert "No workers available" in str(restored)
|
||||||
|
|
||||||
|
def test_error_is_dirty_error_subclass(self):
|
||||||
|
"""DirtyNoWorkersAvailableError is a DirtyError subclass."""
|
||||||
|
error = DirtyNoWorkersAvailableError("app:Class")
|
||||||
|
assert isinstance(error, DirtyError)
|
||||||
|
|
||||||
|
def test_web_app_can_catch_specific_error(self):
|
||||||
|
"""Web app can catch DirtyNoWorkersAvailableError specifically."""
|
||||||
|
def simulate_execute():
|
||||||
|
raise DirtyNoWorkersAvailableError("myapp:HeavyModel")
|
||||||
|
|
||||||
|
# Catch specific error
|
||||||
|
try:
|
||||||
|
simulate_execute()
|
||||||
|
assert False, "Should have raised"
|
||||||
|
except DirtyNoWorkersAvailableError as e:
|
||||||
|
assert e.app_path == "myapp:HeavyModel"
|
||||||
|
|
||||||
|
def test_can_catch_as_base_error(self):
|
||||||
|
"""Can catch DirtyNoWorkersAvailableError as DirtyError."""
|
||||||
|
def simulate_execute():
|
||||||
|
raise DirtyNoWorkersAvailableError("myapp:Model")
|
||||||
|
|
||||||
|
try:
|
||||||
|
simulate_execute()
|
||||||
|
assert False, "Should have raised"
|
||||||
|
except DirtyError as e:
|
||||||
|
# Should catch it as the base class
|
||||||
|
assert hasattr(e, "app_path")
|
||||||
Loading…
x
Reference in New Issue
Block a user