secator 0.15.1__py3-none-any.whl → 0.16.1__py3-none-any.whl

secator/celery.py

@@ -1,8 +1,6 @@
-import gc
 import json
 import logging
 import os
-import uuid
 
 from time import time
 
@@ -33,7 +31,7 @@ logging.basicConfig(
 	handlers=[rich_handler],
 	force=True)
 logging.getLogger('kombu').setLevel(logging.ERROR)
-logging.getLogger('celery').setLevel(logging.INFO if CONFIG.debug.level > 6 else logging.WARNING)
+logging.getLogger('celery').setLevel(logging.DEBUG if 'celery.debug' in CONFIG.debug or 'celery.*' in CONFIG.debug else logging.WARNING)  # noqa: E501
 logger = logging.getLogger(__name__)
 trace.LOG_SUCCESS = "Task %(name)s[%(id)s] succeeded in %(runtime)ss"
 
@@ -169,9 +167,17 @@ def run_scan(self, args=[], kwargs={}):
 @app.task(bind=True)
 def run_command(self, results, name, targets, opts={}):
 	if IN_CELERY_WORKER_PROCESS:
-		opts.update({'print_item': True, 'print_line': True, 'print_cmd': True})
-		# routing_key = self.request.delivery_info['routing_key']
-		# console.print(Info(message=f'Task "{name}" running with routing key "{routing_key}"'))
+		quiet = not CONFIG.cli.worker_command_verbose
+		opts.update({
+			'print_item': True,
+			'print_line': True,
+			'print_cmd': True,
+			'print_target': True,
+			'print_profiles': True,
+			'quiet': quiet
+		})
+		routing_key = self.request.delivery_info['routing_key']
+		debug(f'Task "{name}" running with routing key "{routing_key}"', sub='celery.state')
 
 	# Flatten + dedupe + filter results
 	results = forward_results(results)
@@ -188,11 +194,13 @@ def run_command(self, results, name, targets, opts={}):
 	sync = not IN_CELERY_WORKER_PROCESS
 	task_cls = Task.get_task_class(name)
 	task = task_cls(targets, **opts)
+	chunk_it = task.needs_chunking(sync)
+	task.has_children = chunk_it
 	task.mark_started()
 	update_state(self, task, force=True)
 
 	# Chunk task if needed
-	if task.needs_chunking(sync):
+	if chunk_it:
 		if IN_CELERY_WORKER_PROCESS:
 			console.print(Info(message=f'Task {name} requires chunking, breaking into {len(targets)} tasks'))
 		tasks = break_task(task, opts, results=results)
@@ -200,12 +208,10 @@ def run_command(self, results, name, targets, opts={}):
 		return self.replace(tasks)
 
 	# Update state live
-	[update_state(self, task) for _ in task]
+	for _ in task:
+		update_state(self, task)
 	update_state(self, task, force=True)
 
-	# Garbage collection to save RAM
-	gc.collect()
-
 	return task.results
 
 
@@ -238,8 +244,7 @@ def mark_runner_started(results, runner, enable_hooks=True):
 	if results:
 		runner.results = forward_results(results)
 	runner.enable_hooks = enable_hooks
-	if not runner.dry_run:
-		runner.mark_started()
+	runner.mark_started()
 	return runner.results
 
 
@@ -258,9 +263,9 @@ def mark_runner_completed(results, runner, enable_hooks=True):
 	debug(f'Runner {runner.unique_name} has finished, running mark_completed', sub='celery')
 	results = forward_results(results)
 	runner.enable_hooks = enable_hooks
-	if not runner.dry_run:
-		[runner.add_result(item) for item in results]
-		runner.mark_completed()
+	for item in results:
+		runner.add_result(item, print=False)
+	runner.mark_completed()
 	return runner.results
 
 
@@ -294,7 +299,7 @@ def break_task(task, task_opts, results=[]):
 	)
 
 	# Clone opts
-	opts = task_opts.copy()
+	base_opts = task_opts.copy()
 
 	# Build signatures
 	sigs = []
@@ -302,17 +307,28 @@ def break_task(task, task_opts, results=[]):
 	for ix, chunk in enumerate(chunks):
 		if not isinstance(chunk, list):
 			chunk = [chunk]
-		if len(chunks) > 0:  # add chunk to task opts for tracking chunks exec
-			opts['chunk'] = ix + 1
-			opts['chunk_count'] = len(chunks)
-		task_id = str(uuid.uuid4())
+
+		# Add chunk info to opts
+		opts = base_opts.copy()
+		opts.update({'chunk': ix + 1, 'chunk_count': len(chunks)})
+		debug('', obj={
+			task.unique_name: 'CHUNK',
+			'chunk': f'{ix + 1} / {len(chunks)}',
+			'target_count': len(chunk),
+			'targets': chunk
+		}, sub='celery.state')  # noqa: E501
+
+		# Construct chunked signature
 		opts['has_parent'] = True
 		opts['enable_duplicate_check'] = False
 		opts['results'] = results
-		sig = type(task).si(chunk, **opts).set(queue=type(task).profile, task_id=task_id)
+		if 'targets_' in opts:
+			del opts['targets_']
+		sig = type(task).si(chunk, **opts)
+		task_id = sig.freeze().task_id
 		full_name = f'{task.name}_{ix + 1}'
-		task.add_subtask(task_id, task.name, f'{task.name}_{ix + 1}')
-		info = Info(message=f'Celery chunked task created: {task_id}', _source=full_name, _uuid=str(uuid.uuid4()))
+		task.add_subtask(task_id, task.name, full_name)
+		info = Info(message=f'Celery chunked task created: {task_id}')
 		task.add_result(info)
 		sigs.append(sig)
 

secator/celery_signals.py

@@ -19,116 +19,119 @@ STATE_DIR.mkdir(exist_ok=True, parents=True)
 
 
 def get_lock_file_path():
-    worker_name = os.environ.get("WORKER_NAME", f"unknown_{os.getpid()}")
-    return Path(f"/tmp/celery_worker_{worker_name}.lock")
+	worker_name = os.environ.get("WORKER_NAME", f"unknown_{os.getpid()}")
+	return Path(f"/tmp/celery_worker_{worker_name}.lock")
 
 
 def set_task_running(task_id):
-    """Mark that a task is running in current worker"""
-    with open(get_lock_file_path(), "w") as f:
-        f.write(task_id)
+	"""Mark that a task is running in current worker"""
+	with open(get_lock_file_path(), "w") as f:
+		f.write(task_id)
 
 
 def clear_task_running():
-    """Clear the task running state"""
-    lock_file = get_lock_file_path()
-    if lock_file.exists():
-        lock_file.unlink()
+	"""Clear the task running state"""
+	lock_file = get_lock_file_path()
+	if lock_file.exists():
+		lock_file.unlink()
 
 
 def is_task_running():
-    """Check if a task is currently running"""
-    return get_lock_file_path().exists()
+	"""Check if a task is currently running"""
+	return get_lock_file_path().exists()
 
 
 def kill_worker(parent=False):
-    """Kill current worker using its pid by sending a SIGTERM to Celery master process."""
-    worker_name = os.environ.get('WORKER_NAME', 'unknown')
+	"""Kill current worker using its pid by sending a SIGTERM to Celery master process."""
+	worker_name = os.environ.get('WORKER_NAME', 'unknown')
 
-    # Check if a task is running via the lock file
-    if not is_task_running():
-        pid = os.getppid() if parent else os.getpid()
-        console.print(Info(message=f'Sending SIGTERM to worker {worker_name} with pid {pid}'))
-        os.kill(pid, signal.SIGTERM)
-    else:
-        console.print(Info(message=f'Cancelling worker shutdown of {worker_name} since a task is running'))
+	# Check if a task is running via the lock file
+	if not is_task_running():
+		pid = os.getppid() if parent else os.getpid()
+		console.print(Info(message=f'Sending SIGTERM to worker {worker_name} with pid {pid}'))
+		os.kill(pid, signal.SIGTERM)
+	else:
+		console.print(Info(message=f'Cancelling worker shutdown of {worker_name} since a task is running'))
 
 
 def setup_idle_timer(timeout):
-    """Setup a timer to kill the worker after being idle"""
-    if timeout == -1:
-        return
+	"""Setup a timer to kill the worker after being idle"""
+	if timeout == -1:
+		return
 
-    console.print(Info(message=f'Starting inactivity timer for {timeout} seconds ...'))
-    timer = threading.Timer(timeout, kill_worker)
-    timer.daemon = True  # Make sure timer is killed when worker exits
-    timer.start()
+	console.print(Info(message=f'Starting inactivity timer for {timeout} seconds ...'))
+	timer = threading.Timer(timeout, kill_worker)
+	timer.daemon = True  # Make sure timer is killed when worker exits
+	timer.start()
 
 
 def setup_logging(*args, **kwargs):
-    """Override celery's logging setup to prevent it from altering our settings.
-    github.com/celery/celery/issues/1867
-    """
-    pass
+	"""Override celery's logging setup to prevent it from altering our settings.
+	github.com/celery/celery/issues/1867
+	"""
+	pass
 
 
 def capture_worker_name(sender, instance, **kwargs):
-    os.environ["WORKER_NAME"] = '{0}'.format(sender)
+	os.environ["WORKER_NAME"] = '{0}'.format(sender)
 
 
 def worker_init_handler(**kwargs):
-    if IDLE_TIMEOUT != -1:
-        setup_idle_timer(IDLE_TIMEOUT)
+	if IDLE_TIMEOUT != -1:
+		setup_idle_timer(IDLE_TIMEOUT)
 
 
 def task_prerun_handler(task_id, **kwargs):
-    # Mark that a task is running
-    set_task_running(task_id)
+	# Mark that a task is running
+	set_task_running(task_id)
 
 
 def task_postrun_handler(**kwargs):
-    # Mark that no task is running
-    clear_task_running()
+	# Mark that no task is running
+	clear_task_running()
 
-    # Get sender name from kwargs
-    sender_name = kwargs['sender'].name
+	# Get sender name from kwargs
+	sender_name = kwargs['sender'].name
 
-    if CONFIG.celery.worker_kill_after_task and sender_name.startswith('secator.'):
-        worker_name = os.environ.get('WORKER_NAME', 'unknown')
-        console.print(Info(message=f'Shutdown worker {worker_name} since config celery.worker_kill_after_task is set.'))
-        kill_worker(parent=True)
-        return
+	if CONFIG.celery.worker_kill_after_task and sender_name.startswith('secator.'):
+		worker_name = os.environ.get('WORKER_NAME', 'unknown')
+		console.print(Info(message=f'Shutdown worker {worker_name} since config celery.worker_kill_after_task is set.'))
+		kill_worker(parent=True)
+		return
 
-    # Set up a new idle timer
-    if IDLE_TIMEOUT != -1:
-        console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds'))
-        setup_idle_timer(IDLE_TIMEOUT)
+	# Set up a new idle timer
+	if IDLE_TIMEOUT != -1:
+		console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds'))
+		setup_idle_timer(IDLE_TIMEOUT)
 
 
 def task_revoked_handler(request=None, **kwargs):
-    """Handle revoked tasks by clearing the task running state"""
-    console.print(Info(message='Task was revoked, clearing running state'))
-    clear_task_running()
+	"""Handle revoked tasks by clearing the task running state"""
+	console.print(Info(message='Task was revoked, clearing running state'))
+	clear_task_running()
 
-    # Set up a new idle timer
-    if IDLE_TIMEOUT != -1:
-        console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds after task revocation'))
-        setup_idle_timer(IDLE_TIMEOUT)
+	# Set up a new idle timer
+	if IDLE_TIMEOUT != -1:
+		console.print(Info(message=f'Reset inactivity timer to {IDLE_TIMEOUT} seconds after task revocation'))
+		setup_idle_timer(IDLE_TIMEOUT)
 
 
 def worker_shutdown_handler(**kwargs):
-    """Cleanup lock files when worker shuts down"""
-    lock_file = get_lock_file_path()
-    if lock_file.exists():
-        lock_file.unlink()
+	"""Cleanup lock files when worker shuts down"""
+	lock_file = get_lock_file_path()
+	if lock_file.exists():
+		lock_file.unlink()
 
 
 def setup_handlers():
-    if CONFIG.celery.override_default_logging:
-        signals.setup_logging.connect(setup_logging)
-    signals.celeryd_after_setup.connect(capture_worker_name)
-    signals.task_prerun.connect(task_prerun_handler)
-    signals.task_postrun.connect(task_postrun_handler)
-    signals.task_revoked.connect(task_revoked_handler)
-    signals.worker_ready.connect(worker_init_handler)
-    signals.worker_shutdown.connect(worker_shutdown_handler)
+	if CONFIG.celery.override_default_logging:
+		signals.setup_logging.connect(setup_logging)
+
+	# Register common handlers when either task‐ or idle‐based termination is enabled
+	if CONFIG.celery.worker_kill_after_task or CONFIG.celery.worker_kill_after_idle_seconds != -1:
+		signals.celeryd_after_setup.connect(capture_worker_name)
+		signals.task_postrun.connect(task_postrun_handler)
+		signals.task_prerun.connect(task_prerun_handler)
+		signals.task_revoked.connect(task_revoked_handler)
+		signals.worker_ready.connect(worker_init_handler)
+		signals.worker_shutdown.connect(worker_shutdown_handler)

secator/celery_utils.py

@@ -1,3 +1,5 @@
+import gc
+
 from contextlib import nullcontext
 from time import sleep
 
@@ -5,6 +7,7 @@ import kombu
 import kombu.exceptions
 
 from celery.result import AsyncResult, GroupResult
+from celery.exceptions import TaskRevokedError
 from greenlet import GreenletExit
 from rich.panel import Panel
 from rich.padding import Padding
@@ -24,6 +27,7 @@ class CeleryData(object):
 			result,
 			ids_map={},
 			description=True,
+			revoked=False,
 			refresh_interval=CONFIG.runners.poll_frequency,
 			print_remote_info=True,
 			print_remote_title='Results'
@@ -33,6 +37,7 @@ class CeleryData(object):
 		Args:
 			result (Union[AsyncResult, GroupResult]): Celery result.
 			description (bool): Whether to show task description.
+			revoked (bool): Whether the task was revoked.
 			refresh_interval (int): Refresh interval.
 			print_remote_info (bool): Whether to display live results.
 			print_remote_title (str): Title for the progress panel.
@@ -59,7 +64,7 @@ class CeleryData(object):
 				TextColumn('{task.fields[count]}'),
 				TextColumn('{task.fields[progress]}%'),
 				# TextColumn('\[[bold magenta]{task.fields[id]:<30}[/]]'),  # noqa: W605
-				refresh_per_second=1,
+				auto_refresh=False,
 				transient=False,
 				console=console,
 				# redirect_stderr=True,
@@ -75,7 +80,7 @@ class CeleryData(object):
 				progress_cache = CeleryData.init_progress(progress, ids_map)
 
 			# Get live results and print progress
-			for data in CeleryData.poll(result, ids_map, refresh_interval):
+			for data in CeleryData.poll(result, ids_map, refresh_interval, revoked):
 				for result in data['results']:
 
 					# Add dynamic subtask to ids_map
@@ -93,6 +98,7 @@ class CeleryData(object):
 								'progress': 0
 							}
 					yield result
+					del result
 
 				if print_remote_info:
 					task_id = data['id']
@@ -103,11 +109,17 @@ class CeleryData(object):
 							continue
 					progress_id = progress_cache[task_id]
 					CeleryData.update_progress(progress, progress_id, data)
+					progress.refresh()
+
+				# Garbage collect between polls
+				del data
+				gc.collect()
 
 			# Update all tasks to 100 %
 			if print_remote_info:
 				for progress_id in progress_cache.values():
 					progress.update(progress_id, advance=100)
+				progress.refresh()
 
 	@staticmethod
 	def init_progress(progress, ids_map):
@@ -130,36 +142,22 @@ class CeleryData(object):
 		progress.update(progress_id, **pdata)
 
 	@staticmethod
-	def poll(result, ids_map, refresh_interval):
+	def poll(result, ids_map, refresh_interval, revoked=False):
 		"""Poll Celery subtasks results in real-time. Fetch task metadata and partial results from each task that runs.
 
 		Yields:
 			dict: Subtasks state and results.
 		"""
-		while True:
+		exit_loop = False
+		while not exit_loop:
 			try:
-				main_task = State(
-					task_id=result.id,
-					state=result.state,
-					_source='celery'
-				)
-				debug(f"Main task state: {result.id} - {result.state}", sub='celery.poll', verbose=True)
-				yield {'id': result.id, 'state': result.state, 'results': [main_task]}
-				yield from CeleryData.get_all_data(result, ids_map)
-
-				if result.ready():
+				yield from CeleryData.get_all_data(result, ids_map, revoked=revoked)
+				if result.ready() or revoked:
 					debug('result is ready', sub='celery.poll', id=result.id)
-					main_task = State(
-						task_id=result.id,
-						state=result.state,
-						_source='celery'
-					)
-					debug(f"Final main task state: {result.id} - {result.state}", sub='celery.poll', verbose=True)
-					yield {'id': result.id, 'state': result.state, 'results': [main_task]}
-					yield from CeleryData.get_all_data(result, ids_map)
-					break
+					exit_loop = True
 			except (KeyboardInterrupt, GreenletExit):
 				debug('encounted KeyboardInterrupt or GreenletExit', sub='celery.poll')
+				yield from CeleryData.get_all_data(result, ids_map, revoked=revoked)
 				raise
 			except Exception as e:
 				error = Error.from_exception(e)
@@ -169,7 +167,18 @@ class CeleryData(object):
 				sleep(refresh_interval)
 
 	@staticmethod
-	def get_all_data(result, ids_map):
+	def get_all_data(result, ids_map, revoked=False):
+		main_task = State(
+			task_id=result.id,
+			state='REVOKED' if revoked and result.state == 'PENDING' else result.state,
+			_source='celery'
+		)
+		debug(f"Main task state: {result.id} - {result.state}", sub='celery.poll', verbose=True)
+		yield {'id': result.id, 'state': result.state, 'results': [main_task]}
+		yield from CeleryData.get_tasks_data(ids_map, revoked=revoked)
+
+	@staticmethod
+	def get_tasks_data(ids_map, revoked=False):
 		"""Get Celery results from main result object, AND all subtasks results.
 
 		Yields:
@@ -180,6 +189,8 @@ class CeleryData(object):
 			data = CeleryData.get_task_data(task_id, ids_map)
 			if not data:
 				continue
+			if revoked and data['state'] == 'PENDING':
+				data['state'] = 'REVOKED'
 			debug(
 				'POLL',
 				sub='celery.poll',
@@ -238,13 +249,18 @@ class CeleryData(object):
 		info = res.info
 
 		# Depending on the task state, info will be either an Exception (FAILURE), a list (SUCCESS), or a dict (RUNNING).
-		# - If it's an Exception, it's an unhandled error.
+		# - If it's an Exception, it's a TaskRevokedError or an unhandled error.
 		# - If it's a list, it's the task results.
 		# - If it's a dict, it's the custom user metadata.
 
 		if isinstance(info, Exception):
-			debug('unhandled exception', obj={'msg': str(info), 'tb': traceback_as_string(info)}, sub='celery.data', id=task_id)
-			raise info
+			if isinstance(info, TaskRevokedError):
+				data['results'] = [Error(message='Task was revoked', _source=data['name'])]
+				data['state'] = 'REVOKED'
+				data['ready'] = True
+			else:
+				debug('unhandled exception', obj={'msg': str(info), 'tb': traceback_as_string(info)}, sub='celery.data', id=task_id)
+				raise info
 
 		elif isinstance(info, list):
 			data['results'] = info