mirror of
https://github.com/galaxyproject/galaxy.git
synced 2026-09-24 16:30:27 +08:00
Merge pull request #22189 from mvdbeek/improve-celery-rate-limit-and-concurrnecy-handling
Improve celery rate limit and concurrency handling
This commit is contained in:
Marius van den Beek
@@ -5451,6 +5451,21 @@
|
||||
:Type: float
|
||||
|
||||
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
``celery_user_concurrency_limit``
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
:Description:
|
||||
Maximum number of Celery tasks that can execute concurrently for a
|
||||
single user. If set to 0 (default), no concurrency limit is
|
||||
enforced. When a user exceeds this limit, new tasks are deferred
|
||||
and retried until a slot becomes available. A periodic cleanup
|
||||
task reclaims slots from crashed workers by inspecting active
|
||||
tasks on all workers.
|
||||
:Default: ``0``
|
||||
:Type: int
|
||||
|
||||
|
||||
~~~~~~~~~~~~~~
|
||||
``use_pbkdf2``
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
@@ -262,3 +262,166 @@ This configuration ensures that:
|
||||
2. Exported files remain available for 7 days after generation
|
||||
|
||||
You can monitor Celery task status using [Flower](https://flower.readthedocs.io/en/latest/), a real-time web-based monitoring tool for Celery.
|
||||
|
||||
#### Per-user task rate limiting
|
||||
|
||||
Galaxy supports limiting the rate at which Celery tasks are executed per user. This prevents a single user from monopolizing worker capacity and ensures fair scheduling across all users.
|
||||
|
||||
##### Configuration
|
||||
|
||||
Set `celery_user_rate_limit` in the Galaxy configuration to a non-zero float representing the maximum number of tasks per user per second. For example:
|
||||
|
||||
```yaml
|
||||
celery_user_rate_limit: 0.1
|
||||
```
|
||||
|
||||
This allows each user at most one task execution every 10 seconds. The default value of `0.0` disables rate limiting entirely.
|
||||
|
||||
##### How it works
|
||||
|
||||
Rate limiting is implemented in Celery's `before_start` hook via the `GalaxyTaskBeforeStart` class hierarchy. The mechanism works in two phases:
|
||||
|
||||
1. **Slot reservation (first attempt):** When a task is about to execute for the first time, Galaxy atomically reserves the next available timeslot for that user in the `celery_user_rate_limit` database table. The reserved time is calculated as `max(last_scheduled_time + interval, now)`. If the reserved timeslot is in the future, the task is deferred using `task.retry(countdown=...)` and the reserved time is stored in a Celery message header.
|
||||
|
||||
2. **Execution gating (retries):** On each subsequent retry, the task reads its reserved timeslot from the message header and checks whether the current time has reached it. If not, it retries again with the remaining countdown. Once the timeslot is reached, the task proceeds to execute.
|
||||
|
||||
This two-phase design ensures that:
|
||||
- Each task reserves a slot in the DB exactly once, avoiding cascading delays from re-reservation.
|
||||
- Tasks from different users are scheduled independently — user A's tasks do not delay user B's tasks.
|
||||
- Tasks are retried with `max_retries=None` for rate-limit retries, so they are never lost due to Celery's default retry limit.
|
||||
|
||||
The flow for a single task looks like this:
|
||||
|
||||
```
|
||||
Task arrives (retries=0, no header)
|
||||
→ Atomically reserve timeslot in DB
|
||||
→ Timeslot is now? Execute immediately.
|
||||
→ Timeslot is in the future?
|
||||
→ Store timeslot in message header
|
||||
→ task.retry(countdown=timeslot - now)
|
||||
|
||||
Task wakes up (retry, header present)
|
||||
→ Read reserved timeslot from header
|
||||
→ now >= timeslot? Execute.
|
||||
→ now < timeslot? retry(countdown=remaining)
|
||||
```
|
||||
|
||||
##### Database backends
|
||||
|
||||
Galaxy provides two implementations of the timeslot reservation logic:
|
||||
|
||||
- **PostgreSQL** (`GalaxyTaskBeforeStartUserRateLimitPostgres`): Uses `UPDATE ... RETURNING` with `greatest()` for an atomic, single-statement slot reservation. Falls back to `INSERT ... ON CONFLICT DO UPDATE` (upsert) for the first task by a given user. This is the most efficient implementation.
|
||||
|
||||
- **Standard SQL** (`GalaxyTaskBeforeStartUserRateLimitStandard`): Uses `SELECT ... FOR UPDATE` followed by a separate `UPDATE` (or `INSERT` for new users). This works with SQLite and other databases but requires two statements and explicit locking.
|
||||
|
||||
The correct implementation is selected automatically based on the configured `database_connection`.
|
||||
|
||||
##### Limitations
|
||||
|
||||
- **Rate limiting, not concurrency limiting.** The mechanism controls the *rate* at which tasks are scheduled (tasks per second per user), not how many tasks run concurrently. If a user submits 100 tasks, they will all eventually execute — just spaced apart by the configured interval.
|
||||
- **Tasks without `task_user_id` are not rate-limited.** Only tasks that receive a `task_user_id` keyword argument participate in rate limiting. System tasks and tasks without a user context bypass the check entirely.
|
||||
- **Timeslots are not released on failure.** If a task fails after its timeslot was reserved, that slot is consumed. The next task for the same user will be scheduled after it. This means task failures still "use up" rate-limit capacity.
|
||||
- **Clock precision.** The mechanism relies on `datetime.datetime.now()` on the worker. Clock skew between workers could cause minor scheduling inaccuracies, though this is unlikely to matter in practice.
|
||||
- **No priority or reordering.** Tasks are scheduled in the order they reserve slots (first-come, first-served within a user). There is no mechanism to prioritize certain task types over others for the same user.
|
||||
- **Worker restarts.** If a worker is terminated while holding deferred tasks, Celery's broker (e.g., Redis, RabbitMQ) will redeliver them. The tasks will re-enter the `before_start` hook, read their reserved timeslot from the message header, and continue waiting or execute as appropriate — no slots are lost or duplicated.
|
||||
- **Database overhead.** Each task execution requires one or two queries to the `celery_user_rate_limit` table to reserve a timeslot (one for PostgreSQL's atomic upsert, two for the standard `SELECT FOR UPDATE` + `UPDATE` path). On PostgreSQL at 100 tasks/second this adds ~100 small writes/second; the standard backend doubles that. For most Galaxy deployments (typically fewer than 10 tasks/second) this is negligible. Additionally, tasks that are deferred via `task.retry` re-enter the broker and are redelivered to a worker, adding a small amount of broker traffic proportional to the deferral rate.
|
||||
|
||||
#### Per-user task concurrency limiting
|
||||
|
||||
In addition to rate limiting, Galaxy supports limiting the number of tasks that can execute **concurrently** for a single user. This prevents one user from consuming all available worker capacity.
|
||||
|
||||
##### Configuration
|
||||
|
||||
Set `celery_user_concurrency_limit` in the Galaxy configuration to the maximum number of tasks that can run simultaneously per user:
|
||||
|
||||
```yaml
|
||||
celery_user_concurrency_limit: 5
|
||||
```
|
||||
|
||||
The default value of `0` disables concurrency limiting. This setting can be used independently of or in combination with `celery_user_rate_limit`.
|
||||
|
||||
##### How it works
|
||||
|
||||
Concurrency limiting uses a tracking table (`celery_user_active_task`) that records which tasks are currently executing for each user. The mechanism has three components:
|
||||
|
||||
1. **Admission control (`before_start`):** Before a task executes, the system counts the user's currently active tasks in the tracking table. If the count is at or above the limit, the task is deferred via `task.retry(countdown=5)` with unlimited retries. Otherwise, a tracking row is inserted for this task.
|
||||
|
||||
2. **Cleanup on completion (`after_return`):** When a task finishes (success or failure), its tracking row is deleted from the table. This runs via Celery's `after_return` hook, which fires regardless of whether the task succeeded or failed. Retries do not trigger cleanup — only final completion does.
|
||||
|
||||
3. **Stale row recovery (periodic beat task):** A periodic task (`cleanup_stale_concurrency_slots`) runs every 5 minutes to handle the case where a worker crashes without calling `after_return`. It queries all workers via `celery_app.control.inspect().active()` to get the set of actually-running task IDs, then removes any tracking rows older than 30 minutes whose task ID is not found on any worker.
|
||||
|
||||
The flow for a single task:
|
||||
|
||||
```
|
||||
Task arrives
|
||||
→ Count active tasks for this user in DB
|
||||
→ Count >= limit? → task.retry(countdown=5)
|
||||
→ Count < limit?
|
||||
→ INSERT tracking row (task_id, user_id, started_at)
|
||||
→ Execute task
|
||||
→ Task finishes (success or failure)
|
||||
→ DELETE tracking row
|
||||
|
||||
Periodic cleanup (every 5 min)
|
||||
→ SELECT stale rows (started_at > 30 min ago)
|
||||
→ inspect().active() → get all running task IDs from workers
|
||||
→ DELETE rows where task_id NOT in active set
|
||||
```
|
||||
|
||||
##### Combining rate limiting and concurrency limiting
|
||||
|
||||
When both `celery_user_rate_limit` and `celery_user_concurrency_limit` are set, rate limiting runs first (to schedule the timeslot) and concurrency limiting runs second (to gate execution based on active task count). This means a task must both reach its scheduled timeslot *and* have a concurrency slot available before it can execute.
|
||||
|
||||
##### Limitations
|
||||
|
||||
- **Tasks without `task_user_id` are not limited.** Only tasks that receive a `task_user_id` keyword argument participate in concurrency limiting.
|
||||
- **Worker crash recovery is not instant.** If a worker is killed (SIGKILL, OOM), its tracking rows remain until the periodic cleanup task runs (every 5 minutes by default). During this window, those slots are "leaked" and reduce the user's effective concurrency limit.
|
||||
- **Retry polling interval is fixed.** Deferred tasks retry every 5 seconds. Under heavy load with many deferred tasks, this creates periodic bursts of retry attempts.
|
||||
- **No queue ordering guarantees.** When multiple tasks are waiting for a concurrency slot, the order in which they acquire slots depends on Celery's delivery order and retry timing — not submission order.
|
||||
- **Database overhead.** Each task execution requires an INSERT (on start) and DELETE (on completion) in the tracking table, plus a COUNT query for admission. At 100 tasks/second this adds ~300 small queries/second to the database. For most Galaxy deployments (which typically sustain fewer than 10 tasks/second) this is negligible. Deployments processing hundreds of tasks per second should monitor database connection pool utilization and query latency on the `celery_user_active_task` table.
|
||||
|
||||
##### Administrative operations
|
||||
|
||||
Admins can directly manage the concurrency tracking table and the Celery queue to recover from stuck states or clear backlogs.
|
||||
|
||||
**Clearing leaked concurrency slots manually:**
|
||||
|
||||
If a worker crashes and the periodic cleanup hasn't run yet (or Celery beat is not running), admins can free slots directly:
|
||||
|
||||
```sql
|
||||
-- View all currently tracked active tasks
|
||||
SELECT * FROM celery_user_active_task ORDER BY started_at;
|
||||
|
||||
-- Remove all slots for a specific user (e.g., user_id 42)
|
||||
DELETE FROM celery_user_active_task WHERE user_id = 42;
|
||||
|
||||
-- Remove all stale slots older than 1 hour
|
||||
DELETE FROM celery_user_active_task
|
||||
WHERE started_at < NOW() - INTERVAL '1 hour';
|
||||
|
||||
-- Nuclear option: clear ALL tracking rows (resets all concurrency counters)
|
||||
DELETE FROM celery_user_active_task;
|
||||
```
|
||||
|
||||
After clearing rows, deferred tasks waiting for slots will acquire them on their next retry (within 5 seconds).
|
||||
|
||||
**Purging tasks from the Celery queue:**
|
||||
|
||||
To remove pending (not yet started) tasks from the broker queue:
|
||||
|
||||
```bash
|
||||
# Purge all pending tasks from the default Galaxy queue
|
||||
celery -A galaxy.celery purge -Q galaxy.internal
|
||||
|
||||
# Purge all pending tasks from all queues
|
||||
celery -A galaxy.celery purge
|
||||
|
||||
# Revoke a specific task by ID (prevents it from executing even if already delivered)
|
||||
celery -A galaxy.celery call celery.backend_cleanup # or use the control interface:
|
||||
celery -A galaxy.celery control revoke <task-id>
|
||||
|
||||
# Revoke all pending tasks for inspection first
|
||||
celery -A galaxy.celery inspect reserved
|
||||
```
|
||||
|
||||
Note: `purge` only removes tasks that have not yet been delivered to a worker. Tasks already being executed or waiting in a worker's prefetch buffer require `revoke`. Revoking a task that is mid-execution requires the `--terminate` flag, which sends SIGTERM to the worker process — use with caution.
|
||||
|
||||
+44
-8
@@ -29,7 +29,12 @@ from galaxy.agents.factory import build_registry as build_agent_registry
|
||||
from galaxy.agents.registry import AgentRegistry
|
||||
from galaxy.carbon_emissions import get_carbon_intensity_entry
|
||||
from galaxy.celery.base_task import (
|
||||
GalaxyTaskAfterReturn,
|
||||
GalaxyTaskAfterReturnConcurrencyLimit,
|
||||
GalaxyTaskBeforeStart,
|
||||
GalaxyTaskBeforeStartCombined,
|
||||
GalaxyTaskBeforeStartConcurrencyLimitPostgres,
|
||||
GalaxyTaskBeforeStartConcurrencyLimitStandard,
|
||||
GalaxyTaskBeforeStartUserRateLimitPostgres,
|
||||
GalaxyTaskBeforeStartUserRateLimitStandard,
|
||||
)
|
||||
@@ -709,24 +714,55 @@ class GalaxyManagerApplication(MinimalManagerApp, MinimalGalaxyApplication):
|
||||
|
||||
def _register_celery_galaxy_task_components(self):
|
||||
"""
|
||||
Register subtype class instance to support implementation of a user rate limit for execution of celery tasks.
|
||||
The default supertype class does not enforce a user rate limit. This is the case if the celery_user_rate_limit
|
||||
config param is the default value.
|
||||
Register subtype class instances for user rate limiting and concurrency
|
||||
limiting of celery task executions. Both can be enabled independently
|
||||
or together (combined via GalaxyTaskBeforeStartCombined).
|
||||
"""
|
||||
task_before_start: GalaxyTaskBeforeStart
|
||||
hooks: list[GalaxyTaskBeforeStart] = []
|
||||
|
||||
# Rate limiting hook
|
||||
if self.config.celery_user_rate_limit:
|
||||
if is_postgres(self.config.database_connection): # type: ignore[arg-type]
|
||||
task_before_start = GalaxyTaskBeforeStartUserRateLimitPostgres(
|
||||
self.config.celery_user_rate_limit, self.model.session
|
||||
hooks.append(
|
||||
GalaxyTaskBeforeStartUserRateLimitPostgres(self.config.celery_user_rate_limit, self.model.session)
|
||||
)
|
||||
else:
|
||||
task_before_start = GalaxyTaskBeforeStartUserRateLimitStandard(
|
||||
self.config.celery_user_rate_limit, self.model.session
|
||||
hooks.append(
|
||||
GalaxyTaskBeforeStartUserRateLimitStandard(self.config.celery_user_rate_limit, self.model.session)
|
||||
)
|
||||
|
||||
# Concurrency limiting hook
|
||||
if self.config.celery_user_concurrency_limit:
|
||||
if is_postgres(self.config.database_connection): # type: ignore[arg-type]
|
||||
hooks.append(
|
||||
GalaxyTaskBeforeStartConcurrencyLimitPostgres(
|
||||
self.config.celery_user_concurrency_limit, self.model.session
|
||||
)
|
||||
)
|
||||
else:
|
||||
hooks.append(
|
||||
GalaxyTaskBeforeStartConcurrencyLimitStandard(
|
||||
self.config.celery_user_concurrency_limit, self.model.session
|
||||
)
|
||||
)
|
||||
|
||||
# Register the appropriate before_start hook
|
||||
if len(hooks) > 1:
|
||||
task_before_start: GalaxyTaskBeforeStart = GalaxyTaskBeforeStartCombined(*hooks)
|
||||
elif len(hooks) == 1:
|
||||
task_before_start = hooks[0]
|
||||
else:
|
||||
task_before_start = GalaxyTaskBeforeStart()
|
||||
self._register_singleton(GalaxyTaskBeforeStart, task_before_start)
|
||||
|
||||
# Register after_return hook for concurrency tracking cleanup
|
||||
task_after_return: GalaxyTaskAfterReturn
|
||||
if self.config.celery_user_concurrency_limit:
|
||||
task_after_return = GalaxyTaskAfterReturnConcurrencyLimit(self.model.session)
|
||||
else:
|
||||
task_after_return = GalaxyTaskAfterReturn()
|
||||
self._register_singleton(GalaxyTaskAfterReturn, task_after_return)
|
||||
|
||||
def _configure_tool_shed_registry(self) -> None:
|
||||
# Set up the tool sheds registry
|
||||
if os.path.isfile(self.config.tool_sheds_config_file):
|
||||
|
||||
@@ -23,7 +23,10 @@ from celery.signals import (
|
||||
)
|
||||
from kombu import serialization
|
||||
|
||||
from galaxy.celery.base_task import GalaxyTaskBeforeStart
|
||||
from galaxy.celery.base_task import (
|
||||
GalaxyTaskAfterReturn,
|
||||
GalaxyTaskBeforeStart,
|
||||
)
|
||||
from galaxy.config import Configuration
|
||||
from galaxy.main_config import find_config
|
||||
from galaxy.util import ExecutionTimer
|
||||
@@ -71,8 +74,8 @@ class GalaxyCelery(Celery):
|
||||
|
||||
class GalaxyTask(Task):
|
||||
"""
|
||||
Custom celery task used to limit number of tasks executions per user
|
||||
per second.
|
||||
Custom celery task used to enforce per-user rate limits and
|
||||
concurrency limits on task executions.
|
||||
"""
|
||||
|
||||
def before_start(self, task_id, args, kwargs):
|
||||
@@ -83,6 +86,17 @@ class GalaxyTask(Task):
|
||||
assert app
|
||||
app[GalaxyTaskBeforeStart](self, task_id, args, kwargs)
|
||||
|
||||
def after_return(self, status, retval, task_id, args, kwargs, einfo):
|
||||
"""
|
||||
Called after task returns (success, failure, revoked, or retry).
|
||||
Used to clean up concurrency tracking rows.
|
||||
"""
|
||||
if status == "RETRY":
|
||||
return # Don't clean up on retry — the task will run again
|
||||
app = get_galaxy_app()
|
||||
if app:
|
||||
app[GalaxyTaskAfterReturn](self, task_id, args, kwargs)
|
||||
|
||||
|
||||
def set_thread_app(app):
|
||||
APP_LOCAL.app = app
|
||||
@@ -251,6 +265,10 @@ def setup_periodic_tasks(config, celery_app):
|
||||
if config.vault_token_renewal_interval:
|
||||
schedule_task("renew_vault_token", config.vault_token_renewal_interval)
|
||||
|
||||
if config.celery_user_concurrency_limit:
|
||||
# Run cleanup every 5 minutes (300 seconds)
|
||||
schedule_task("cleanup_stale_concurrency_slots", 300)
|
||||
|
||||
if beat_schedule:
|
||||
celery_app.conf.beat_schedule = beat_schedule
|
||||
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
import datetime
|
||||
import logging
|
||||
from abc import abstractmethod
|
||||
from typing import cast
|
||||
|
||||
from celery import Task
|
||||
from sqlalchemy import (
|
||||
bindparam,
|
||||
delete,
|
||||
func,
|
||||
insert,
|
||||
select,
|
||||
@@ -15,7 +17,16 @@ from sqlalchemy.engine import CursorResult
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import scoped_session
|
||||
|
||||
from galaxy.model import CeleryUserRateLimit
|
||||
from galaxy.model import (
|
||||
CeleryUserActiveTask,
|
||||
CeleryUserRateLimit,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
HEADER_SCHEDULED_TIME = "_gxy_rate_limit_scheduled_time"
|
||||
HEADER_CONCURRENCY_TRACKED = "_gxy_concurrency_tracked"
|
||||
CONCURRENCY_RETRY_COUNTDOWN_SECS = 5.0
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStart:
|
||||
@@ -31,6 +42,16 @@ class GalaxyTaskBeforeStart:
|
||||
pass
|
||||
|
||||
|
||||
class GalaxyTaskAfterReturn:
|
||||
"""
|
||||
Hook called after a task returns (success, failure, or retry).
|
||||
Base class is a no-op; subclasses implement concurrency tracking cleanup.
|
||||
"""
|
||||
|
||||
def __call__(self, task: Task, task_id, args, kwargs):
|
||||
pass
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStartUserRateLimit(GalaxyTaskBeforeStart):
|
||||
"""
|
||||
Used when we wish to enforce a user rate limit based on
|
||||
@@ -42,6 +63,11 @@ class GalaxyTaskBeforeStartUserRateLimit(GalaxyTaskBeforeStart):
|
||||
by doing a task.retry.
|
||||
If the last scheduled execution was far enough in the past
|
||||
then we allow the task to run immediately.
|
||||
|
||||
The reserved timeslot is stored in a message header so that on
|
||||
retry the task can verify it has reached its scheduled time
|
||||
without re-reserving a new slot. This ensures tasks are never
|
||||
lost and will keep retrying until their timeslot arrives.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -56,17 +82,39 @@ class GalaxyTaskBeforeStartUserRateLimit(GalaxyTaskBeforeStart):
|
||||
self.ga_scoped_session = ga_scoped_session
|
||||
|
||||
def __call__(self, task: Task, task_id, args, kwargs):
|
||||
if task.request.retries > 0:
|
||||
return
|
||||
usr = kwargs.get("task_user_id")
|
||||
if not usr:
|
||||
return
|
||||
now = datetime.datetime.now()
|
||||
sa_session = self.ga_scoped_session
|
||||
next_scheduled_time = self.calculate_task_start_time(usr, sa_session, self.task_exec_countdown_secs, now)
|
||||
if next_scheduled_time > now:
|
||||
count_down = next_scheduled_time - now
|
||||
task.retry(countdown=count_down.total_seconds())
|
||||
|
||||
# Check if this task already has a reserved timeslot from a previous attempt
|
||||
headers = task.request.headers or {}
|
||||
reserved_time_str = headers.get(HEADER_SCHEDULED_TIME)
|
||||
|
||||
if reserved_time_str:
|
||||
# Retry path: verify we've reached our reserved timeslot
|
||||
reserved_time = datetime.datetime.fromisoformat(reserved_time_str)
|
||||
if now >= reserved_time:
|
||||
return # Timeslot reached, proceed with execution
|
||||
# Not yet time — retry with remaining countdown
|
||||
remaining = (reserved_time - now).total_seconds()
|
||||
task.retry(
|
||||
countdown=remaining,
|
||||
max_retries=None,
|
||||
headers={HEADER_SCHEDULED_TIME: reserved_time_str},
|
||||
)
|
||||
else:
|
||||
# First attempt: reserve a timeslot atomically in the DB
|
||||
sa_session = self.ga_scoped_session
|
||||
next_scheduled_time = self.calculate_task_start_time(usr, sa_session, self.task_exec_countdown_secs, now)
|
||||
if next_scheduled_time > now:
|
||||
count_down = next_scheduled_time - now
|
||||
task.retry(
|
||||
countdown=count_down.total_seconds(),
|
||||
max_retries=None,
|
||||
headers={HEADER_SCHEDULED_TIME: next_scheduled_time.isoformat()},
|
||||
)
|
||||
# else: scheduled time is now or in the past, execute immediately
|
||||
|
||||
@abstractmethod
|
||||
def calculate_task_start_time(
|
||||
@@ -157,3 +205,129 @@ class GalaxyTaskBeforeStartUserRateLimitStandard(GalaxyTaskBeforeStartUserRateLi
|
||||
raise Exception(f"Failed to update a celery_user_rate_limit row for user id {user_id}")
|
||||
sa_session.commit()
|
||||
return sched_time
|
||||
|
||||
|
||||
# --- Per-user concurrency limiting ---
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStartConcurrencyLimit(GalaxyTaskBeforeStart):
|
||||
"""
|
||||
Enforces a per-user concurrency limit on Celery task execution.
|
||||
Before a task starts, checks if the user already has the maximum
|
||||
number of tasks running. If so, defers execution via task.retry().
|
||||
|
||||
On successful admission, inserts a tracking row into celery_user_active_task
|
||||
and sets a header so after_return knows to clean it up.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_concurrent: int,
|
||||
ga_scoped_session: scoped_session,
|
||||
):
|
||||
self.max_concurrent = max_concurrent
|
||||
self.ga_scoped_session = ga_scoped_session
|
||||
|
||||
def __call__(self, task: Task, task_id, args, kwargs):
|
||||
usr = kwargs.get("task_user_id")
|
||||
if not usr:
|
||||
return
|
||||
|
||||
headers = task.request.headers or {}
|
||||
|
||||
# If this task was already admitted (retry after concurrency admission),
|
||||
# don't re-check concurrency — it already has a tracking row.
|
||||
if headers.get(HEADER_CONCURRENCY_TRACKED):
|
||||
return
|
||||
|
||||
sa_session = self.ga_scoped_session
|
||||
now = datetime.datetime.now()
|
||||
|
||||
# Count currently active tasks for this user
|
||||
active_count = self._get_active_count(usr, sa_session)
|
||||
|
||||
if active_count >= self.max_concurrent:
|
||||
# User is at capacity — defer this task
|
||||
sa_session.commit()
|
||||
task.retry(
|
||||
countdown=CONCURRENCY_RETRY_COUNTDOWN_SECS,
|
||||
max_retries=None,
|
||||
)
|
||||
return
|
||||
|
||||
# Admit this task: insert tracking row
|
||||
try:
|
||||
sa_session.execute(
|
||||
insert(CeleryUserActiveTask).values(
|
||||
task_id=str(task_id),
|
||||
user_id=usr,
|
||||
started_at=now,
|
||||
)
|
||||
)
|
||||
sa_session.commit()
|
||||
except IntegrityError:
|
||||
# Task ID already tracked (e.g., redelivery) — that's fine
|
||||
sa_session.rollback()
|
||||
|
||||
@abstractmethod
|
||||
def _get_active_count(self, user_id: int, sa_session: scoped_session) -> int: ...
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStartConcurrencyLimitPostgres(GalaxyTaskBeforeStartConcurrencyLimit):
|
||||
"""Postgres-optimized concurrency check using SELECT COUNT with row-level advisory awareness."""
|
||||
|
||||
def _get_active_count(self, user_id: int, sa_session: scoped_session) -> int:
|
||||
count = sa_session.scalar(
|
||||
select(func.count()).select_from(CeleryUserActiveTask).where(CeleryUserActiveTask.user_id == user_id)
|
||||
)
|
||||
return count or 0
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStartConcurrencyLimitStandard(GalaxyTaskBeforeStartConcurrencyLimit):
|
||||
"""Standard SQL concurrency check."""
|
||||
|
||||
def _get_active_count(self, user_id: int, sa_session: scoped_session) -> int:
|
||||
count = sa_session.scalar(
|
||||
select(func.count()).select_from(CeleryUserActiveTask).where(CeleryUserActiveTask.user_id == user_id)
|
||||
)
|
||||
return count or 0
|
||||
|
||||
|
||||
class GalaxyTaskAfterReturnConcurrencyLimit(GalaxyTaskAfterReturn):
|
||||
"""
|
||||
Cleans up the concurrency tracking row after a task completes
|
||||
(regardless of success or failure).
|
||||
"""
|
||||
|
||||
def __init__(self, ga_scoped_session: scoped_session):
|
||||
self.ga_scoped_session = ga_scoped_session
|
||||
|
||||
def __call__(self, task: Task, task_id, args, kwargs):
|
||||
usr = kwargs.get("task_user_id")
|
||||
if not usr:
|
||||
return
|
||||
|
||||
sa_session = self.ga_scoped_session
|
||||
try:
|
||||
sa_session.execute(delete(CeleryUserActiveTask).where(CeleryUserActiveTask.task_id == str(task_id)))
|
||||
sa_session.commit()
|
||||
except Exception:
|
||||
log.exception(f"Failed to remove concurrency tracking row for task {task_id}")
|
||||
sa_session.rollback()
|
||||
|
||||
|
||||
# --- Combined before_start that chains rate limit + concurrency limit ---
|
||||
|
||||
|
||||
class GalaxyTaskBeforeStartCombined(GalaxyTaskBeforeStart):
|
||||
"""
|
||||
Chains multiple before_start hooks. Rate limiting runs first
|
||||
(to schedule the timeslot), then concurrency limiting (to gate execution).
|
||||
"""
|
||||
|
||||
def __init__(self, *hooks: GalaxyTaskBeforeStart):
|
||||
self.hooks = hooks
|
||||
|
||||
def __call__(self, task: Task, task_id, args, kwargs):
|
||||
for hook in self.hooks:
|
||||
hook(task, task_id, args, kwargs)
|
||||
|
||||
@@ -12,6 +12,7 @@ from typing import (
|
||||
|
||||
from sqlalchemy import (
|
||||
and_,
|
||||
delete,
|
||||
exists,
|
||||
false,
|
||||
select,
|
||||
@@ -689,3 +690,48 @@ def execute_workflow_completion_hook(
|
||||
log.info(f"Successfully executed hook '{hook_name}' for invocation {invocation_id}")
|
||||
else:
|
||||
log.error(f"Failed to execute hook '{hook_name}' for invocation {invocation_id}")
|
||||
|
||||
|
||||
@galaxy_task(action="clean up stale concurrency tracking rows")
|
||||
def cleanup_stale_concurrency_slots(
|
||||
session: galaxy_scoped_session,
|
||||
stale_threshold_minutes: int = 30,
|
||||
):
|
||||
"""
|
||||
Periodic task that reclaims concurrency slots from tasks that are
|
||||
no longer running (e.g., due to worker crashes). Queries all workers
|
||||
for their active tasks and removes tracking rows for any task that
|
||||
is no longer executing on any worker.
|
||||
"""
|
||||
now = datetime.datetime.now()
|
||||
threshold = now - datetime.timedelta(minutes=stale_threshold_minutes)
|
||||
|
||||
# Only consider rows older than the threshold — recent tasks are likely still running
|
||||
stale_rows = (
|
||||
session.execute(
|
||||
select(model.CeleryUserActiveTask.task_id).where(model.CeleryUserActiveTask.started_at < threshold)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
|
||||
if not stale_rows:
|
||||
return
|
||||
|
||||
# Ask all workers what they're currently running
|
||||
try:
|
||||
active_response = celery_app.control.inspect().active() or {}
|
||||
except Exception:
|
||||
log.warning("Failed to inspect active tasks on workers; skipping stale cleanup")
|
||||
return
|
||||
|
||||
active_task_ids = {t["id"] for tasks in active_response.values() for t in tasks}
|
||||
|
||||
# Remove tracking rows for tasks that are NOT on any worker
|
||||
stale_task_ids = [tid for tid in stale_rows if tid not in active_task_ids]
|
||||
if stale_task_ids:
|
||||
session.execute(
|
||||
delete(model.CeleryUserActiveTask).where(model.CeleryUserActiveTask.task_id.in_(stale_task_ids))
|
||||
)
|
||||
session.commit()
|
||||
log.info(f"Cleaned up {len(stale_task_ids)} stale concurrency tracking rows")
|
||||
|
||||
@@ -2968,6 +2968,14 @@ galaxy:
|
||||
# executed per user per second.
|
||||
#celery_user_rate_limit: 0.0
|
||||
|
||||
# Maximum number of Celery tasks that can execute concurrently for a
|
||||
# single user. If set to 0 (default), no concurrency limit is
|
||||
# enforced. When a user exceeds this limit, new tasks are deferred and
|
||||
# retried until a slot becomes available. A periodic cleanup task
|
||||
# reclaims slots from crashed workers by inspecting active tasks on
|
||||
# all workers.
|
||||
#celery_user_concurrency_limit: 0
|
||||
|
||||
# Allow disabling pbkdf2 hashing of passwords for legacy situations.
|
||||
# This should normally be left enabled unless there is a specific
|
||||
# reason to disable it.
|
||||
|
||||
@@ -4039,6 +4039,18 @@ mapping:
|
||||
If set to a non-0 value, upper limit on number of
|
||||
tasks that can be executed per user per second.
|
||||
|
||||
celery_user_concurrency_limit:
|
||||
type: int
|
||||
default: 0
|
||||
required: false
|
||||
desc: |
|
||||
Maximum number of Celery tasks that can execute concurrently
|
||||
for a single user. If set to 0 (default), no concurrency
|
||||
limit is enforced. When a user exceeds this limit, new tasks
|
||||
are deferred and retried until a slot becomes available.
|
||||
A periodic cleanup task reclaims slots from crashed workers
|
||||
by inspecting active tasks on all workers.
|
||||
|
||||
use_pbkdf2:
|
||||
type: bool
|
||||
default: true
|
||||
|
||||
@@ -12633,6 +12633,20 @@ class CeleryUserRateLimit(Base):
|
||||
)
|
||||
|
||||
|
||||
class CeleryUserActiveTask(Base):
|
||||
"""
|
||||
Tracks actively executing Celery tasks per user for concurrency limiting.
|
||||
Each row represents a task that has started but not yet completed.
|
||||
A periodic cleanup task removes stale rows from crashed workers.
|
||||
"""
|
||||
|
||||
__tablename__ = "celery_user_active_task"
|
||||
|
||||
task_id: Mapped[str] = mapped_column(String(255), primary_key=True)
|
||||
user_id: Mapped[int] = mapped_column(ForeignKey("galaxy_user.id", ondelete="CASCADE"), index=True)
|
||||
started_at: Mapped[datetime]
|
||||
|
||||
|
||||
class UserCredentials(Base):
|
||||
"""
|
||||
Represents a credential associated with a user for a specific service.
|
||||
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
"""create celery_user_active_task table
|
||||
|
||||
Revision ID: f5e9e4bca542
|
||||
Revises: 566b691307a5
|
||||
Create Date: 2026-03-19 10:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from galaxy.model.migrations.util import (
|
||||
create_table,
|
||||
drop_table,
|
||||
)
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "f5e9e4bca542"
|
||||
down_revision = "566b691307a5"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
create_table(
|
||||
"celery_user_active_task",
|
||||
sa.Column("task_id", sa.String(255), primary_key=True),
|
||||
sa.Column("user_id", sa.Integer, sa.ForeignKey("galaxy_user.id", ondelete="CASCADE"), index=True),
|
||||
sa.Column("started_at", sa.DateTime, nullable=False),
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
drop_table("celery_user_active_task")
|
||||
@@ -0,0 +1,273 @@
|
||||
import datetime
|
||||
import tempfile
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
|
||||
from celery.result import AsyncResult
|
||||
from sqlalchemy import (
|
||||
select,
|
||||
text,
|
||||
)
|
||||
|
||||
from galaxy.celery import galaxy_task
|
||||
from galaxy.model import CeleryUserActiveTask
|
||||
from galaxy.model.database_utils import sqlalchemy_engine
|
||||
from galaxy.model.scoped_session import galaxy_scoped_session
|
||||
from galaxy_test.driver.driver_util import init_database
|
||||
from galaxy_test.driver.integration_util import (
|
||||
IntegrationTestCase,
|
||||
skip_unless_postgres,
|
||||
)
|
||||
|
||||
|
||||
@galaxy_task(action="sleep for concurrency testing")
|
||||
def mock_sleep_task(
|
||||
session: galaxy_scoped_session,
|
||||
sleep_seconds: float = 2.0,
|
||||
task_user_id: Optional[int] = None,
|
||||
):
|
||||
"""Task that sleeps for a configurable duration, used to test concurrency limits."""
|
||||
time.sleep(sleep_seconds)
|
||||
return task_user_id
|
||||
|
||||
|
||||
@galaxy_task
|
||||
def mock_fast_task(task_user_id: int):
|
||||
"""Instant task used to verify cleanup after completion."""
|
||||
return task_user_id
|
||||
|
||||
|
||||
@lru_cache
|
||||
def sqlite_url():
|
||||
path = tempfile.NamedTemporaryFile().name
|
||||
dburl = f"sqlite:///{path}"
|
||||
init_database(dburl)
|
||||
return dburl
|
||||
|
||||
|
||||
@lru_cache
|
||||
def setup_users(dburl: str, num_users: int = 3):
|
||||
"""
|
||||
Setup test users in galaxy_user table with user id's starting from 2.
|
||||
"""
|
||||
expected_user_ids = list(range(2, num_users + 2))
|
||||
with sqlalchemy_engine(dburl) as engine:
|
||||
with engine.begin() as conn:
|
||||
found_user_ids = conn.scalars(
|
||||
text("select id from galaxy_user where id between 2 and :high"), {"high": num_users + 1}
|
||||
).all()
|
||||
if len(expected_user_ids) > len(found_user_ids):
|
||||
user_ids_to_add = set(expected_user_ids).difference(found_user_ids)
|
||||
for user_id in user_ids_to_add:
|
||||
conn.execute(
|
||||
text("insert into galaxy_user(id, active, email, password) values (:id, :active, :email, :pw)"),
|
||||
[{"id": user_id, "active": True, "email": f"e{user_id}", "pw": "p"}],
|
||||
)
|
||||
|
||||
|
||||
class TestCeleryUserConcurrencyLimitIntegration(IntegrationTestCase):
|
||||
"""
|
||||
Base class for per-user concurrency limiting tests.
|
||||
Does not define test_* methods directly — subclasses call _test_* helpers.
|
||||
"""
|
||||
|
||||
_concurrency_limit = 2
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
|
||||
def _get_active_task_count(self, user_id: int) -> int:
|
||||
"""Query the tracking table for active tasks for a given user."""
|
||||
sa_session = self._app.model.session()
|
||||
try:
|
||||
count = len(
|
||||
sa_session.execute(select(CeleryUserActiveTask.task_id).where(CeleryUserActiveTask.user_id == user_id))
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
return count
|
||||
finally:
|
||||
sa_session.close()
|
||||
|
||||
def _test_concurrency_limit_enforced(self):
|
||||
"""
|
||||
Submit more tasks than the concurrency limit for a single user.
|
||||
With concurrency_limit=2 and 4 tasks each sleeping 2s, the tasks
|
||||
should take ~4s total (2 batches of 2). Without limits, they'd
|
||||
all run in ~2s.
|
||||
"""
|
||||
user_id = 2
|
||||
num_tasks = 4
|
||||
sleep_seconds = 2.0
|
||||
|
||||
start = datetime.datetime.now(datetime.timezone.utc)
|
||||
results: list[AsyncResult] = []
|
||||
for _ in range(num_tasks):
|
||||
results.append(mock_sleep_task.delay(sleep_seconds=sleep_seconds, task_user_id=user_id))
|
||||
|
||||
# Collect all results
|
||||
for result in results:
|
||||
val = result.get(timeout=120)
|
||||
assert val == user_id
|
||||
|
||||
elapsed = (datetime.datetime.now(datetime.timezone.utc) - start).total_seconds()
|
||||
# With concurrency=2, 4 tasks sleeping 2s each should take ~4s
|
||||
# (2 execute, finish, then the next 2 execute)
|
||||
expected_min = sleep_seconds * (num_tasks / self._concurrency_limit) - 1
|
||||
# Allow generous upper bound for scheduling overhead
|
||||
expected_max = sleep_seconds * (num_tasks / self._concurrency_limit) + 15
|
||||
assert elapsed >= expected_min, (
|
||||
f"Tasks completed too fast ({elapsed:.1f}s < {expected_min:.1f}s), "
|
||||
f"concurrency limit may not be enforced"
|
||||
)
|
||||
assert elapsed <= expected_max, f"Tasks took too long ({elapsed:.1f}s > {expected_max:.1f}s)"
|
||||
|
||||
def _test_different_users_independent(self):
|
||||
"""
|
||||
Tasks from different users should run independently.
|
||||
User A and User B each submit 2 tasks (at concurrency limit).
|
||||
All 4 tasks should complete in ~2s (parallel across users),
|
||||
not 4s (if users shared a limit).
|
||||
"""
|
||||
user_a = 2
|
||||
user_b = 3
|
||||
sleep_seconds = 2.0
|
||||
|
||||
start = datetime.datetime.now(datetime.timezone.utc)
|
||||
results: list[AsyncResult] = []
|
||||
# Submit concurrency_limit tasks for each user
|
||||
for user_id in [user_a, user_b]:
|
||||
for _ in range(self._concurrency_limit):
|
||||
results.append(mock_sleep_task.delay(sleep_seconds=sleep_seconds, task_user_id=user_id))
|
||||
|
||||
for result in results:
|
||||
val = result.get(timeout=120)
|
||||
assert val in (user_a, user_b)
|
||||
|
||||
elapsed = (datetime.datetime.now(datetime.timezone.utc) - start).total_seconds()
|
||||
# Both users run their tasks concurrently — should take ~2s, not ~4s
|
||||
expected_max = sleep_seconds + 15 # generous overhead
|
||||
assert elapsed <= expected_max, (
|
||||
f"Cross-user tasks took too long ({elapsed:.1f}s > {expected_max:.1f}s), "
|
||||
f"users may be sharing a concurrency limit"
|
||||
)
|
||||
|
||||
def _test_tracking_rows_cleaned_up(self):
|
||||
"""
|
||||
After tasks complete, their tracking rows should be removed
|
||||
from celery_user_active_task.
|
||||
"""
|
||||
user_id = 2
|
||||
results = []
|
||||
for _ in range(3):
|
||||
results.append(mock_fast_task.delay(task_user_id=user_id))
|
||||
|
||||
# Wait for all tasks to complete
|
||||
for result in results:
|
||||
result.get(timeout=60)
|
||||
|
||||
# Give a moment for after_return to fire
|
||||
time.sleep(1)
|
||||
|
||||
active_count = self._get_active_task_count(user_id)
|
||||
assert active_count == 0, f"Expected 0 active tracking rows after completion, found {active_count}"
|
||||
|
||||
def _test_tasks_without_user_id_bypass_limit(self):
|
||||
"""
|
||||
Tasks that don't provide task_user_id should bypass concurrency limiting.
|
||||
"""
|
||||
# Submit tasks without task_user_id — they should run immediately
|
||||
results = []
|
||||
for _ in range(5):
|
||||
results.append(mock_fast_task.delay(task_user_id=0))
|
||||
|
||||
# If concurrency limiting incorrectly applied, these would queue up
|
||||
for result in results:
|
||||
result.get(timeout=30)
|
||||
|
||||
|
||||
@skip_unless_postgres()
|
||||
class TestCeleryUserConcurrencyLimitPostgres(TestCeleryUserConcurrencyLimitIntegration):
|
||||
_concurrency_limit = 2
|
||||
|
||||
@classmethod
|
||||
def handle_galaxy_config_kwds(cls, config):
|
||||
super().handle_galaxy_config_kwds(config)
|
||||
config["celery_user_concurrency_limit"] = cls._concurrency_limit
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
dburl = self._app.config.database_connection
|
||||
setup_users(dburl, num_users=3)
|
||||
|
||||
def test_concurrency_limit_enforced(self):
|
||||
self._test_concurrency_limit_enforced()
|
||||
|
||||
def test_different_users_independent(self):
|
||||
self._test_different_users_independent()
|
||||
|
||||
def test_tracking_rows_cleaned_up(self):
|
||||
self._test_tracking_rows_cleaned_up()
|
||||
|
||||
def test_tasks_without_user_id_bypass_limit(self):
|
||||
self._test_tasks_without_user_id_bypass_limit()
|
||||
|
||||
|
||||
class TestCeleryUserConcurrencyLimitSqlite(TestCeleryUserConcurrencyLimitIntegration):
|
||||
_concurrency_limit = 2
|
||||
|
||||
@classmethod
|
||||
def handle_galaxy_config_kwds(cls, config):
|
||||
super().handle_galaxy_config_kwds(config)
|
||||
config["check_migrate_databases"] = False
|
||||
config["database_connection"] = sqlite_url()
|
||||
if config.get("database_engine_option_pool_size"):
|
||||
config.pop("database_engine_option_pool_size")
|
||||
if config.get("database_engine_option_max_overflow"):
|
||||
config.pop("database_engine_option_max_overflow")
|
||||
config["celery_user_concurrency_limit"] = cls._concurrency_limit
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
dburl = self._app.config.database_connection
|
||||
setup_users(dburl, num_users=3)
|
||||
|
||||
def test_concurrency_limit_enforced(self):
|
||||
self._test_concurrency_limit_enforced()
|
||||
|
||||
def test_different_users_independent(self):
|
||||
self._test_different_users_independent()
|
||||
|
||||
def test_tracking_rows_cleaned_up(self):
|
||||
self._test_tracking_rows_cleaned_up()
|
||||
|
||||
def test_tasks_without_user_id_bypass_limit(self):
|
||||
self._test_tasks_without_user_id_bypass_limit()
|
||||
|
||||
|
||||
class TestCeleryUserConcurrencyLimitDisabled(IntegrationTestCase):
|
||||
"""Test that with concurrency_limit=0 (disabled), tasks run without restriction."""
|
||||
|
||||
@classmethod
|
||||
def handle_galaxy_config_kwds(cls, config):
|
||||
super().handle_galaxy_config_kwds(config)
|
||||
config["check_migrate_databases"] = False
|
||||
config["database_connection"] = sqlite_url()
|
||||
if config.get("database_engine_option_pool_size"):
|
||||
config.pop("database_engine_option_pool_size")
|
||||
if config.get("database_engine_option_max_overflow"):
|
||||
config.pop("database_engine_option_max_overflow")
|
||||
|
||||
def test_all_tasks_run_without_restriction(self):
|
||||
"""With no limit, tasks should complete without concurrency deferral."""
|
||||
user_id = 2
|
||||
|
||||
results = []
|
||||
for _ in range(5):
|
||||
results.append(mock_fast_task.delay(task_user_id=user_id))
|
||||
|
||||
# All tasks should complete — none should be stuck waiting for a slot
|
||||
for result in results:
|
||||
val = result.get(timeout=30)
|
||||
assert val == user_id
|
||||
Reference in New Issue
Block a user