def rate_limit(state, task_name, rate_limit, **kwargs):
"""Tell worker(s) to modify the rate limit for a task by type.
See Also:
:attr:`celery.task.base.Task.rate_limit`.
Arguments:
task_name (str): Type of task to set rate limit for.
rate_limit (int, str): New rate limit.
"""
try:
rate(rate_limit)
except ValueError as exc:
return nok('Invalid rate limit string: {0!r}'.format(exc))
try:
state.app.tasks[task_name].rate_limit = rate_limit
except KeyError:
logger.error('Rate limit attempt for unknown task %s',
task_name,
exc_info=True)
return nok('unknown task')
state.consumer.reset_rate_limits()
if not rate_limit:
logger.info('Rate limits disabled for tasks of type %s', task_name)
return ok('rate limit disabled successfully')
logger.info('New rate limit for tasks of type %s: %s.', task_name,
rate_limit)
return ok('new rate limit set successfully')
def rate_limit(state, task_name, rate_limit, **kwargs):
"""Tell worker(s) to modify the rate limit for a task by type.
See Also:
:attr:`celery.task.base.Task.rate_limit`.
Arguments:
task_name (str): Type of task to set rate limit for.
rate_limit (int, str): New rate limit.
"""
# pylint: disable=redefined-outer-name
# XXX Note that this redefines `terminate`:
# Outside of this scope that is a function.
try:
rate(rate_limit)
except ValueError as exc:
return nok('Invalid rate limit string: {0!r}'.format(exc))
try:
state.app.tasks[task_name].rate_limit = rate_limit
except KeyError:
logger.error('Rate limit attempt for unknown task %s',
task_name, exc_info=True)
return nok('unknown task')
state.consumer.reset_rate_limits()
if not rate_limit:
logger.info('Rate limits disabled for tasks of type %s', task_name)
return ok('rate limit disabled successfully')
logger.info('New rate limit for tasks of type %s: %s.',
task_name, rate_limit)
return ok('new rate limit set successfully')
def rate_limit(state, task_name, rate_limit, **kwargs):
"""Tell worker(s) to modify the rate limit for a task by type.
See Also:
:attr:`celery.task.base.Task.rate_limit`.
Arguments:
task_name (str): Type of task to set rate limit for.
rate_limit (int, str): New rate limit.
"""
# pylint: disable=redefined-outer-name
# XXX Note that this redefines `terminate`:
# Outside of this scope that is a function.
try:
rate(rate_limit)
except ValueError as exc:
return nok(f'Invalid rate limit string: {exc!r}')
try:
state.app.tasks[task_name].rate_limit = rate_limit
except KeyError:
logger.error('Rate limit attempt for unknown task %s',
task_name,
exc_info=True)
return nok('unknown task')
state.consumer.reset_rate_limits()
if not rate_limit:
logger.info('Rate limits disabled for tasks of type %s', task_name)
return ok('rate limit disabled successfully')
logger.info('New rate limit for tasks of type %s: %s.', task_name,
rate_limit)
return ok('new rate limit set successfully')
def _context(self,
sig,
rate_limits=True,
events=True,
utc=True,
limit=None):
assert sig.type.Strategy
assert sig.type.Request
reserved = Mock()
consumer = Mock()
consumer.task_buckets = defaultdict(lambda: None)
if limit:
bucket = TokenBucket(rate(limit), capacity=1)
consumer.task_buckets[sig.task] = bucket
consumer.controller.state.revoked = set()
consumer.disable_rate_limits = not rate_limits
consumer.event_dispatcher.enabled = events
s = sig.type.start_strategy(self.app, consumer, task_reserved=reserved)
assert s
message = self.task_message_from_sig(
self.app,
sig,
utc=utc,
TaskMessage=self.get_message_class(),
)
message = self.prepare_message(message)
yield self.Context(sig, s, reserved, consumer, message)
def __init__(self, state, freq=1.0, maxrate=None,
cleanup_freq=3600.0, timer=None, app=None):
self.app = app_or_default(app)
self.state = state
self.freq = freq
self.cleanup_freq = cleanup_freq
self.timer = timer or self.timer or Timer()
self.logger = logger
self.maxrate = maxrate and TokenBucket(rate(maxrate))
def test_conversion(self):
self.assertEqual(rate(999), 999)
self.assertEqual(rate(7.5), 7.5)
self.assertEqual(rate('2.5/s'), 2.5)
self.assertEqual(rate('1456/s'), 1456)
self.assertEqual(rate('100/m'),
100 / 60.0)
self.assertEqual(rate('10/h'),
10 / 60.0 / 60.0)
for zero in (0, None, '0', '0/m', '0/h', '0/s', '0.0/s'):
self.assertEqual(rate(zero), 0)
def __init__(self,
state,
freq=1.0,
maxrate=None,
cleanup_freq=3600.0,
timer=None,
app=None):
self.app = app_or_default(app)
self.state = state
self.freq = freq
self.cleanup_freq = cleanup_freq
self.timer = timer or self.timer or Timer()
self.logger = logger
self.maxrate = maxrate and TokenBucket(rate(maxrate))
def _context(self, sig,
rate_limits=True, events=True, utc=True, limit=None):
assert sig.type.Strategy
reserved = Mock()
consumer = Mock()
consumer.task_buckets = defaultdict(lambda: None)
if limit:
bucket = TokenBucket(rate(limit), capacity=1)
consumer.task_buckets[sig.task] = bucket
consumer.controller.state.revoked = set()
consumer.disable_rate_limits = not rate_limits
consumer.event_dispatcher.enabled = events
s = sig.type.start_strategy(self.app, consumer, task_reserved=reserved)
assert s
message = self.task_message_from_sig(
self.app, sig, utc=utc, TaskMessage=self.get_message_class(),
)
message = self.prepare_message(message)
yield self.Context(sig, s, reserved, consumer, message)
def bucket_for_task(self, type):
limit = rate(getattr(type, 'rate_limit', None))
return TokenBucket(limit, capacity=1) if limit else None
def bucket_for_task(self, type):
limit = rate(getattr(type, 'rate_limit', None))
return TokenBucket(limit, capacity=1) if limit else None
def test_rate_limit_string(s, expected):
assert rate(s) == expected
def test_rate_limit_string(s, expected):
assert rate(s) == expected