abstract-block-dumper 0.0.5__py3-none-any.whl

abstract_block_dumper/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.0.5'
+__version_tuple__ = version_tuple = (0, 0, 5)
+
+__commit_id__ = commit_id = None

abstract_block_dumper/admin.py

@@ -0,0 +1,73 @@
+from django.contrib import admin
+
+from abstract_block_dumper.models import TaskAttempt
+
+
+@admin.register(TaskAttempt)
+class TaskAttemptAdmin(admin.ModelAdmin):
+    list_display = [
+        "executable_path",
+        "block_number",
+        "status",
+    ]
+    list_filter = [
+        "status",
+        "executable_path",
+    ]
+    search_fields = ["celery_task_id", "block_number"]
+    readonly_fields = [
+        "block_number",
+        "executable_path",
+        "args_json",
+        "status",
+        # Execution fields
+        "celery_task_id",
+        "execution_result",
+        # Attempts & Retry fields
+        "last_attempted_at",
+        "attempt_count",
+        "next_retry_at",
+        "created_at",
+        "updated_at",
+    ]
+    fieldsets = (
+        (
+            None,
+            {
+                "fields": (
+                    "block_number",
+                    "executable_path",
+                    "args_json",
+                )
+            },
+        ),
+        (
+            "Task Execution",
+            {
+                "fields": (
+                    "status",
+                    "celery_task_id",
+                    "execution_result",
+                )
+            },
+        ),
+        (
+            "Retry Information",
+            {
+                "fields": (
+                    "last_attempted_at",
+                    "attempt_count",
+                    "next_retry_at",
+                )
+            },
+        ),
+        (
+            "Timestamps",
+            {
+                "fields": (
+                    "created_at",
+                    "updated_at",
+                )
+            },
+        ),
+    )

abstract_block_dumper/apps.py

@@ -0,0 +1,7 @@
+from django.apps import AppConfig
+
+
+class AbstractBlockDumperConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "abstract_block_dumper"
+    verbose_name = "Abstract Block Dumper"

abstract_block_dumper/dal/django_dal.py

@@ -0,0 +1,150 @@
+from datetime import timedelta
+from typing import Any
+
+from django.conf import settings
+from django.db import transaction
+from django.db.models.query import QuerySet
+from django.utils import timezone
+
+import abstract_block_dumper.models as abd_models
+import abstract_block_dumper.services.utils as abd_utils
+
+
+def get_ready_to_retry_attempts() -> QuerySet[abd_models.TaskAttempt]:
+    return abd_models.TaskAttempt.objects.filter(
+        next_retry_at__isnull=False,
+        next_retry_at__lte=timezone.now(),
+        attempt_count__lt=abd_utils.get_max_attempt_limit(),
+    ).exclude(
+        status=abd_models.TaskAttempt.Status.SUCCESS,
+    )
+
+
+def executed_block_numbers(executable_path: str, args_json: str, from_block: int, to_block: int) -> set[int]:
+    block_numbers = abd_models.TaskAttempt.objects.filter(
+        executable_path=executable_path,
+        args_json=args_json,
+        block_number__gte=from_block,
+        block_number__lt=to_block,
+        status=abd_models.TaskAttempt.Status.SUCCESS,
+    ).values_list("block_number", flat=True)
+    return set(block_numbers)
+
+
+def reset_to_pending(task: abd_models.TaskAttempt) -> None:
+    task.celery_task_id = None
+    task.status = abd_models.TaskAttempt.Status.PENDING
+    task.save()
+
+
+def revert_to_failed(task: abd_models.TaskAttempt) -> None:
+    task.status = abd_models.TaskAttempt.Status.FAILED
+    task.save()
+
+
+def get_recent_phantom_tasks() -> QuerySet[abd_models.TaskAttempt]:
+    """
+    Get tasks marked as SUCCESS but never actually started.
+
+    Only clean up recent phantom tasks to avoid deleting legitimate external successes
+    """
+    return abd_models.TaskAttempt.objects.filter(
+        status=abd_models.TaskAttempt.Status.SUCCESS,
+        last_attempted_at__isnull=True,
+        celery_task_id__isnull=True,  # Additional safety check
+        created_at__gte=timezone.now() - timedelta(hours=1),  # Only recent tasks
+    )
+
+
+def task_can_retry(task: abd_models.TaskAttempt) -> bool:
+    blocked_statuses = {task.Status.SUCCESS, task.Status.RUNNING}
+    return task.status not in blocked_statuses and task.attempt_count < abd_utils.get_max_attempt_limit()
+
+
+def task_mark_as_started(task: abd_models.TaskAttempt, celery_task_id: str) -> None:
+    task.celery_task_id = celery_task_id
+    task.status = abd_models.TaskAttempt.Status.RUNNING
+    task.last_attempted_at = timezone.now()
+    task.save()
+
+
+def task_mark_as_success(task: abd_models.TaskAttempt, result_data: dict) -> None:
+    task.status = task.Status.SUCCESS
+    task.execution_result = result_data
+    task.last_attempted_at = timezone.now()
+    task.next_retry_at = None
+    task.save()
+
+
+def task_mark_as_failed(task) -> None:
+    DEFAULT_BLOCK_TASK_RETRY_BACKOFF = 1
+    MAX_RETRY_DELAY_MINUTES = 1440  # 24 hours max delay
+
+    task.status = task.Status.FAILED
+    task.last_attempted_at = timezone.now()
+    task.attempt_count += 1
+
+    if task_can_retry(task):
+        base_retry_backoff = getattr(settings, "BLOCK_TASK_RETRY_BACKOFF", DEFAULT_BLOCK_TASK_RETRY_BACKOFF)
+        max_delay_minutes = getattr(settings, "BLOCK_TASK_MAX_RETRY_DELAY_MINUTES", MAX_RETRY_DELAY_MINUTES)
+
+        # Calculate exponential backoff with bounds checking
+        backoff_minutes = base_retry_backoff**task.attempt_count
+        backoff_minutes = min(backoff_minutes, max_delay_minutes)
+
+        task.next_retry_at = timezone.now() + timedelta(minutes=backoff_minutes)
+    else:
+        task.next_retry_at = None
+    task.save()
+
+
+def task_schedule_to_retry(task):
+    task.status = abd_models.TaskAttempt.Status.PENDING
+    task.save()
+
+
+def task_create_or_get_pending(
+    block_number: int,
+    executable_path: str,
+    args: dict[str, Any] | None = None,
+) -> tuple[abd_models.TaskAttempt, bool]:
+    """
+    Create or get a pending task attempt.
+    Returns (task, created) where created indicates if a new task was created.
+
+    For failed tasks that can retry:
+    - If next_retry_at is in the future, leave task as FAILED (will be picked up by scheduler)
+    - If next_retry_at is in the past or None, reset to PENDING for immediate execution
+    """
+    if args is None:
+        args = {}
+
+    args_json = abd_utils.serialize_args(args)
+
+    with transaction.atomic():
+        task, created = abd_models.TaskAttempt.objects.get_or_create(
+            block_number=block_number,
+            executable_path=executable_path,
+            args_json=args_json,
+            defaults={"status": abd_models.TaskAttempt.Status.PENDING},
+        )
+
+        # Don't modify tasks that are already in a terminal or active state
+        active_state = {abd_models.TaskAttempt.Status.SUCCESS, abd_models.TaskAttempt.Status.RUNNING}
+        if created or task.status in active_state:
+            return task, created
+
+        # For failed tasks that can retry, only reset to PENDING if retry time has passed
+        if task.status == abd_models.TaskAttempt.Status.FAILED and task_can_retry(task):
+            now = timezone.now()
+            if task.next_retry_at is None or task.next_retry_at <= now:
+                task.status = abd_models.TaskAttempt.Status.PENDING
+                task.save()
+    return task, created
+
+
+def get_the_latest_executed_block_number() -> int | None:
+    qs = abd_models.TaskAttempt.objects.order_by("-block_number").first()
+    if qs:
+        return qs.block_number
+    return None

abstract_block_dumper/dal/memory_registry.py

@@ -0,0 +1,105 @@
+import abc
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+import structlog
+from celery import Task
+
+from abstract_block_dumper.exceptions import ConditionEvaluationError
+
+logger = structlog.getLogger(__name__)
+
+
+@dataclass
+class RegistryItem:
+    condition: Callable[..., bool]
+    function: Task
+    args: list[dict[str, Any]] | None = None
+    backfilling_lookback: int | None = None
+    celery_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    def match_condition(self, block_number: int, **kwargs) -> bool:
+        """
+        Check if condition matches for given block and arguments
+        """
+        try:
+            return self.condition(block_number, **kwargs)
+        except Exception as e:
+            logger.error(
+                "Condition evaluation failed",
+                condition=self.function.__name__,
+                block_number=block_number,
+                exc_info=True,
+            )
+            raise ConditionEvaluationError(f"Failed to evaluate condition: {e}") from e
+
+    def get_execution_args(self) -> list[dict[str, Any]]:
+        """
+        Get list of argument sets for execution
+        """
+        return self.args or [{}]
+
+    @property
+    def executable_path(self) -> str:
+        """
+        Get the importable path to the function.
+        """
+        if hasattr(self.function, "name") and self.function.name is not None:
+            return self.function.name
+
+        return ".".join([self.function.__module__, self.function.__name__])
+
+    def requires_backfilling(self) -> bool:
+        """
+        Check if this item requires backfilling.
+        """
+        return self.backfilling_lookback is not None
+
+
+class BaseRegistry(abc.ABC):
+    @abc.abstractmethod
+    def register_item(self, item: RegistryItem) -> None:
+        pass
+
+    @abc.abstractmethod
+    def get_functions(self) -> list[RegistryItem]:
+        pass
+
+    @abc.abstractmethod
+    def clear(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def get_by_executable_path(self, executable_path: str) -> RegistryItem | None:
+        pass
+
+
+class MemoryRegistry(BaseRegistry):
+    _functions: list[RegistryItem] = []
+
+    def register_item(self, item: RegistryItem) -> None:
+        self._functions.append(item)
+        logger.info(
+            "Registered function",
+            function_name=item.function.__name__,
+            executable_path=item.executable_path,
+            args=item.args,
+            backfilling_lookback=item.backfilling_lookback,
+        )
+
+    def get_functions(self) -> list[RegistryItem]:
+        return self._functions
+
+    def clear(self) -> None:
+        self._functions = []
+
+    def get_by_executable_path(self, executable_path: str) -> RegistryItem:
+        for registry_item in self.get_functions():
+            if registry_item.executable_path == executable_path:
+                return registry_item
+        # TODO: Improve this
+        raise Exception("Function Not Found")
+
+
+task_registry = MemoryRegistry()

abstract_block_dumper/decorators.py

@@ -0,0 +1,211 @@
+from collections.abc import Callable
+from typing import Any, cast
+
+import structlog
+from celery import Task, shared_task
+from django.db import OperationalError, transaction
+
+import abstract_block_dumper.dal.django_dal as abd_dal
+import abstract_block_dumper.services.utils as abd_utils
+from abstract_block_dumper.dal.memory_registry import RegistryItem, task_registry
+from abstract_block_dumper.exceptions import CeleryTaskLocked
+from abstract_block_dumper.models import TaskAttempt
+
+logger = structlog.get_logger(__name__)
+
+
+def schedule_retry(task_attempt: TaskAttempt) -> None:
+    """
+    Schedule a retry for a failed task by calling the decorated Celery task directly.
+
+    Task must already be in FAILED state with next_retry_at set by mark_failed()
+    """
+
+    if not task_attempt.next_retry_at:
+        logger.error(
+            "Cannot schedule retry without next_retry_at",
+            task_id=task_attempt.id,
+            block_number=task_attempt.block_number,
+            executable_path=task_attempt.executable_path,
+        )
+
+    if task_attempt.status != TaskAttempt.Status.FAILED:
+        logger.warning(
+            "Attempted to schedule retry for non-failed task",
+            task_id=task_attempt.id,
+            status=task_attempt.status,
+        )
+        return
+
+    logger.info(
+        "Scheduling retry",
+        task_id=task_attempt.id,
+        attempt_count=task_attempt.attempt_count,
+        next_retry_at=task_attempt.next_retry_at,
+    )
+
+    abd_dal.task_schedule_to_retry(task_attempt)
+
+    celery_task = task_registry.get_by_executable_path(task_attempt.executable_path)
+    if not celery_task:
+        logger.error(
+            "Cannot schedule retry - task not found in registry",
+            executable_path=task_attempt.executable_path,
+        )
+        return
+
+    celery_task.function.apply_async(
+        kwargs={
+            "block_number": task_attempt.block_number,
+            **task_attempt.args_dict,
+        },
+        eta=task_attempt.next_retry_at,
+    )
+
+
+def _celery_task_wrapper(func, block_number: int, **kwargs) -> dict[str, Any] | None:
+    executable_path = abd_utils.get_executable_path(func)
+
+    with transaction.atomic():
+        try:
+            task_attempt = TaskAttempt.objects.select_for_update(nowait=True).get(
+                block_number=block_number,
+                executable_path=executable_path,
+                args_json=abd_utils.serialize_args(kwargs),
+            )
+        except TaskAttempt.DoesNotExist:
+            logger.warning(
+                "TaskAttempt not found - task may have been canceled directly",
+                block_number=block_number,
+                executable_path=executable_path,
+            )
+            raise CeleryTaskLocked("TaskAttempt not found - task may have been canceled directly")
+        except OperationalError as e:
+            logger.info(
+                "Task already being processed by another worker",
+                block_number=block_number,
+                executable_path=executable_path,
+                operational_error=str(e),
+            )
+            raise CeleryTaskLocked("Task already being processed by another worker")
+
+        if task_attempt.status != TaskAttempt.Status.PENDING:
+            logger.info(
+                "Task already processed",
+                task_id=task_attempt.id,
+                status=task_attempt.status,
+            )
+            return None
+
+        abd_dal.task_mark_as_started(task_attempt, abd_utils.get_current_celery_task_id())
+
+        # Start task execution
+        try:
+            execution_kwargs = {"block_number": block_number, **kwargs}
+            logger.info(
+                "Starting task execution",
+                task_id=task_attempt.id,
+                block_number=block_number,
+                executable_path=executable_path,
+                celery_task_id=task_attempt.celery_task_id,
+                execution_kwargs=execution_kwargs,
+            )
+
+            result = func(**execution_kwargs)
+
+            abd_dal.task_mark_as_success(task_attempt, result)
+
+            logger.info("Task completed successfully", task_id=task_attempt.id)
+            return {"result": result}
+        except Exception as e:
+            logger.error(
+                "Task execution failed",
+                task_id=task_attempt.id,
+                error_type=type(e).__name__,
+                exc_info=True,
+            )
+            abd_dal.task_mark_as_failed(task_attempt)
+
+    # Schedule retry after transaction commits:
+    if abd_dal.task_can_retry(task_attempt):
+        try:
+            schedule_retry(task_attempt)
+        except Exception:
+            logger.error(
+                "Failed to schedule retry",
+                task_id=task_attempt.id,
+                exc_info=True,
+            )
+    return None
+
+
+def block_task(
+    condition: Callable[..., bool],
+    args: list[dict[str, Any]] | None = None,
+    backfilling_lookback: int | None = None,
+    celery_kwargs: dict[str, Any] | None = None,
+) -> Callable[..., Any]:
+    """
+    Decorator for registering block tasks.
+
+    Args:
+        condition: Lambda function that determines when to execute
+        args: List of argument dictionaries for multi-execution
+        backfilling_lookback: Number of blocks to backfill
+        celery_kwargs: Additional Celery task parameters
+
+    Examples:
+        @block_task(
+            condition=lambda bn: bn % 100 == 0
+        )
+        def simple_task(block_number: int):
+            pass
+
+        @block_task(
+            condition=lambda bn, netuid: bn + netuid % 100 == 0,
+            args=[{"netuid": 3}, {"netuid": 22}],
+            backfilling_lookback=300,
+            celery_kwargs={"queue": "high-priority"}
+        )
+        def multi_netuid_task(block_number: int, netuid: int):
+            pass
+
+    """
+
+    def decorator(func: Callable[..., Any]) -> Any:
+        if not callable(condition):
+            raise ValueError("condition must be a callable.")
+
+        # Celery task wrapper
+        def shared_celery_task(block_number: int, **kwargs) -> None | Any:
+            """
+            Wrapper that handles TaskAttempt tracking and executed the original
+            function
+
+            This entire wrapper becomes a Celery task.
+            """
+            return _celery_task_wrapper(func, block_number, **kwargs)
+
+        # Wrap with celery shared_task
+        celery_task = shared_task(
+            name=abd_utils.get_executable_path(func),
+            bind=False,
+            **celery_kwargs or {},
+        )(shared_celery_task)
+
+        # Store original function referefence for introspection
+        celery_task._original_func = func
+
+        # Register the Celery task
+        task_registry.register_item(
+            RegistryItem(
+                condition=condition,
+                function=cast(Task, celery_task),
+                args=args,
+                backfilling_lookback=backfilling_lookback,
+                celery_kwargs=celery_kwargs or {},
+            )
+        )
+        return celery_task
+
+    return decorator

abstract_block_dumper/discovery.py

@@ -0,0 +1,24 @@
+import importlib
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+def ensure_modules_loaded() -> None:
+    """
+    Ensure common tasks modules are imported to trigger @block_task registration.
+
+    @block_task must be loaded, otherwise it won't be registered.
+    """
+    from django.apps import apps
+
+    for app_config in apps.get_app_configs():
+        for module_suffix in ["tasks", "block_tasks"]:
+            try:
+                importlib.import_module(f"{app_config.name}.{module_suffix}")
+            except ModuleNotFoundError:
+                continue
+            except ImportError as e:
+                logger.warning(f"Failed to import {app_config.name}.{module_suffix}: {e}")
+                continue

abstract_block_dumper/exceptions.py

@@ -0,0 +1,16 @@
+class AbstractBlockDumperError(Exception):
+    """Base exception for all Abstract Block Dumper errors."""
+
+    pass
+
+
+class ConditionEvaluationError(AbstractBlockDumperError):
+    """Condition failed to evaluate."""
+
+    pass
+
+
+class CeleryTaskLocked(Exception):
+    """Celery task execution is locked"""
+
+    pass

abstract_block_dumper/management/commands/block_tasks.py

@@ -0,0 +1,19 @@
+from django.core.management.base import BaseCommand
+
+from abstract_block_dumper.dal.memory_registry import task_registry
+from abstract_block_dumper.discovery import ensure_modules_loaded
+from abstract_block_dumper.services.scheduler import task_scheduler_factory
+
+
+class Command(BaseCommand):
+    help = "Run the block scheduler daemon."
+
+    def handle(self, *args, **options) -> None:
+        self.stdout.write("Syncing decorated functions...")
+        ensure_modules_loaded()
+        functions_counter = len(task_registry.get_functions())
+        self.stdout.write(self.style.SUCCESS(f"Synced {functions_counter} functions"))
+
+        scheduler = task_scheduler_factory()
+        self.stdout.write("Starting block scheduler...")
+        scheduler.start()