ergon-framework-python 0.1.0__py3-none-any.whl

ergon/task/policies.py

@@ -0,0 +1,229 @@
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field, field_validator
+
+# =====================================================================
+#   INTERNAL NORMALIZERS (framework-level helpers)
+# =====================================================================
+
+
+def _normalize_optional(v):
+    """
+    Normalize optional env-driven values.
+
+    Accepts:
+      - None
+      - ""
+      - "none" / "null"
+      - numeric strings
+
+    Lets Pydantic handle final coercion.
+    """
+    if v is None:
+        return None
+
+    if isinstance(v, str):
+        v = v.strip()
+        if v == "" or v.lower() in {"none", "null"}:
+            return None
+
+    return v
+
+
+def _normalize_bool(v):
+    if isinstance(v, str):
+        return v.strip().lower() in {"1", "true", "yes", "on"}
+    return v
+
+
+# =====================================================================
+#   SHARED MODELS
+# =====================================================================
+
+
+class ConcurrencyPolicy(BaseModel):
+    value: int = Field(default=1, ge=1)
+    headroom: int = Field(default=0, ge=0)
+    min: int = Field(default=1, ge=1)
+    max: int = Field(default=1, ge=1)
+
+    @field_validator("value", "headroom", "min", "max", mode="before")
+    @classmethod
+    def _normalize_ints(cls, v):
+        return _normalize_optional(v)
+
+
+class BatchIntervalPolicy(BaseModel):
+    backoff: float = Field(default=0.0, ge=0.0)
+    backoff_multiplier: float = Field(default=1.0, ge=0.0)
+    backoff_cap: float = Field(default=0.0, ge=0.0)
+    interval: float = Field(default=0.0, ge=0.0)
+
+    @field_validator(
+        "backoff",
+        "backoff_multiplier",
+        "backoff_cap",
+        "interval",
+        mode="before",
+    )
+    @classmethod
+    def _normalize_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+class BatchPolicy(BaseModel):
+    size: int = Field(default=1, ge=1)
+    min_size: int = Field(default=1, ge=1)
+    max_size: int = Field(default=1, ge=1)
+    interval: BatchIntervalPolicy = Field(default_factory=BatchIntervalPolicy)
+
+    @field_validator("size", "min_size", "max_size", "interval", mode="before")
+    @classmethod
+    def _normalize_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+class RetryPolicy(BaseModel):
+    max_attempts: int = Field(default=1, ge=1)
+    timeout: Optional[float] = Field(default=None, ge=0)
+    backoff: float = Field(default=0.0, ge=0.0)
+    backoff_multiplier: float = Field(default=0.0, ge=1.0)
+    backoff_cap: float = Field(default=0.0, ge=0.0)
+
+    @field_validator(
+        "max_attempts",
+        "timeout",
+        "backoff",
+        "backoff_multiplier",
+        "backoff_cap",
+        mode="before",
+    )
+    @classmethod
+    def _normalize_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+class TransactionRuntimePolicy(BaseModel):
+    timeout: Optional[float] = Field(default=60.0, ge=0)
+
+    @field_validator("timeout", mode="before")
+    @classmethod
+    def _normalize_optional_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+# =====================================================================
+#   CONSUMER STEP POLICIES
+# =====================================================================
+
+
+class EmptyFetchPolicy(BaseModel):
+    backoff: float = Field(default=0.0, ge=0.0)
+    backoff_multiplier: float = Field(default=1.0, ge=0.0)
+    backoff_cap: float = Field(default=0.0, ge=0.0)
+    interval: float = Field(default=0.0, ge=0.0)
+
+    @field_validator(
+        "backoff",
+        "backoff_multiplier",
+        "backoff_cap",
+        "interval",
+        mode="before",
+    )
+    @classmethod
+    def _normalize_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+class FetchPolicy(BaseModel):
+    retry: RetryPolicy = Field(default_factory=RetryPolicy)
+    batch: BatchPolicy = Field(default_factory=BatchPolicy)
+    empty: EmptyFetchPolicy = Field(default_factory=EmptyFetchPolicy)
+    connector_name: Optional[str] = None
+    extra: dict = Field(default_factory=dict)
+
+
+class ProcessPolicy(BaseModel):
+    retry: RetryPolicy = Field(default_factory=RetryPolicy)
+
+
+class SuccessPolicy(BaseModel):
+    retry: RetryPolicy = Field(default_factory=RetryPolicy)
+
+
+class ExceptionPolicy(BaseModel):
+    retry: RetryPolicy = Field(default_factory=RetryPolicy)
+
+
+# =====================================================================
+#   CONSUMER LOOP POLICIES
+# =====================================================================
+
+
+class ConsumerLoopPolicy(BaseModel):
+    concurrency: ConcurrencyPolicy = Field(default_factory=ConcurrencyPolicy)
+    timeout: Optional[float] = Field(default=None, ge=0)
+    limit: Optional[int] = Field(default=None, ge=0)
+    streaming: bool = Field(default=False)
+
+    @field_validator("timeout", "limit", mode="before")
+    @classmethod
+    def _normalize_optional_numbers(cls, v):
+        return _normalize_optional(v)
+
+    @field_validator("streaming", mode="before")
+    @classmethod
+    def _normalize_streaming(cls, v):
+        return _normalize_bool(v)
+
+
+# =====================================================================
+#   CONSUMER POLICY
+# =====================================================================
+
+
+class ConsumerPolicy(BaseModel):
+    name: Optional[str] = None
+    loop: ConsumerLoopPolicy = Field(default_factory=ConsumerLoopPolicy)
+    fetch: FetchPolicy = Field(default_factory=FetchPolicy)
+    transaction_runtime: TransactionRuntimePolicy = Field(default_factory=TransactionRuntimePolicy)
+    process: ProcessPolicy = Field(default_factory=ProcessPolicy)
+    success: SuccessPolicy = Field(default_factory=SuccessPolicy)
+    exception: ExceptionPolicy = Field(default_factory=ExceptionPolicy)
+
+
+# =====================================================================
+#   PRODUCER POLICIES
+# =====================================================================
+
+
+class PreparePolicy(BaseModel):
+    retry: RetryPolicy = Field(default_factory=RetryPolicy)
+
+
+class ProducerLoopPolicy(BaseModel):
+    concurrency: ConcurrencyPolicy = Field(default_factory=ConcurrencyPolicy)
+    batch: BatchPolicy = Field(default_factory=BatchPolicy)
+    timeout: Optional[float] = Field(default=None, ge=0)
+    limit: Optional[int] = Field(default=None, ge=0)
+
+    @field_validator("timeout", "limit", mode="before")
+    @classmethod
+    def _normalize_optional_numbers(cls, v):
+        return _normalize_optional(v)
+
+
+class ProducerPolicy(BaseModel):
+    name: Optional[str] = None
+    loop: ProducerLoopPolicy = Field(default_factory=ProducerLoopPolicy)
+    transaction_runtime: TransactionRuntimePolicy = Field(default_factory=TransactionRuntimePolicy)
+    prepare: PreparePolicy = Field(default_factory=PreparePolicy)
+    success: SuccessPolicy = Field(default_factory=SuccessPolicy)
+    exception: ExceptionPolicy = Field(default_factory=ExceptionPolicy)
+
+
+# =====================================================================
+#   UNION
+# =====================================================================
+
+ExecutionPolicy = Union[ConsumerPolicy, ProducerPolicy]

ergon/task/runner.py

@@ -0,0 +1,386 @@
+import asyncio
+import os
+import signal
+import threading
+import traceback
+import uuid
+from concurrent.futures import ProcessPoolExecutor
+from datetime import datetime
+from enum import IntEnum
+from typing import Any, Literal
+
+from ..connector import Transaction
+from ..telemetry import logging, metrics, tracing
+from .base import (
+    BaseAsyncTask,
+    BaseTask,
+    TaskConfig,
+    TaskExecMetadata,
+)
+
+# =============================================================
+# EXIT CODES (POSIX-ALIGNED)
+# =============================================================
+
+
+class ExitCode(IntEnum):
+    SUCCESS = 0
+    ERROR = 1
+    CONFIG_ERROR = 2
+
+    SIGINT = 130  # 128 + SIGINT(2)
+    SIGTERM = 143  # 128 + SIGTERM(15)
+
+
+# =============================================================
+# SHUTDOWN STATE (PROCESS-LOCAL)
+# =============================================================
+
+_shutdown_event = threading.Event()
+_shutdown_signal: int | None = None
+
+
+def _signal_handler(signum, frame):
+    global _shutdown_signal
+    _shutdown_signal = signum
+    _shutdown_event.set()
+
+
+def _install_signal_handlers():
+    signal.signal(signal.SIGINT, _signal_handler)
+    signal.signal(signal.SIGTERM, _signal_handler)
+
+
+def is_shutdown_requested() -> bool:
+    return _shutdown_event.is_set()
+
+
+def get_shutdown_exit_code() -> ExitCode:
+    if _shutdown_signal == signal.SIGINT:
+        return ExitCode.SIGINT
+    if _shutdown_signal == signal.SIGTERM:
+        return ExitCode.SIGTERM
+    return ExitCode.ERROR
+
+
+# =============================================================
+# TELEMETRY INITIALIZATION
+# =============================================================
+
+
+def __init_telemetry(config: TaskConfig, task: object, task_exec_metadata: dict[str, Any]):
+    if config.logging is not None:
+        logging._apply_logging_config(cfg=config.logging, task=task, metadata=task_exec_metadata)
+
+    if config.tracing is not None:
+        tracing._apply_tracing_config(cfg=config.tracing, metadata=task_exec_metadata)
+
+    if config.metrics is not None:
+        metrics._apply_metrics_config(cfg=config.metrics, metadata=task_exec_metadata)
+
+
+# =============================================================
+# ASYNC TRANSACTION EXECUTION
+# =============================================================
+
+
+async def __run_transaction_async(
+    instance: BaseAsyncTask,
+    policy: str,
+    transaction: Transaction | None = None,
+    transaction_id: str | None = None,
+):
+    policy_obj = next((p for p in instance.policies if p.name == policy), None)
+    if not policy_obj:
+        raise ValueError(f"Policy '{policy}' not found")
+
+    if not transaction and not transaction_id:
+        raise ValueError("Either transaction or transaction_id must be provided")
+
+    if transaction_id:
+        conn = instance._resolve_connector(policy_obj.fetch.connector_name)  # type: ignore[attr-defined]
+        transaction = await conn.fetch_transaction_by_id_async(transaction_id)
+
+    success, result = await instance._start_processing(transaction, policy_obj)  # type: ignore[attr-defined]
+    if not success:
+        raise result
+    return result
+
+
+# =============================================================
+# ASYNC TASK EXECUTION
+# =============================================================
+
+
+async def __run_task_async(
+    config: TaskConfig,
+    mode: Literal["task", "transaction"] = "task",
+    *args,
+    **kwargs,
+):
+    if not issubclass(config.task, BaseAsyncTask):  # type: ignore[arg-type]
+        raise ValueError(f"Invalid async task: {config.task}")
+
+    worker_id = kwargs.pop("worker_id", None)
+
+    task_exec_metadata = TaskExecMetadata(
+        task_name=config.name,
+        execution_id=str(uuid.uuid4()),
+        execution_start_time=datetime.now().isoformat(),
+        pid=os.getpid(),
+        worker_id=worker_id,
+    ).model_dump()
+
+    __init_telemetry(config, task=config.task, task_exec_metadata=task_exec_metadata)
+    tracer = tracing.get_tracer(f"task.{config.name}")
+
+    instance = None
+
+    try:
+        with tracer.start_as_current_span(  # type: ignore[attr-defined]
+            f"{config.task.__name__}.run",
+            attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+        ):
+            connectors = {}
+            for name, cfg in config.connectors.items():
+                conn = cfg.connector(*cfg.args, **cfg.kwargs)
+                if hasattr(conn, "init_async"):
+                    await conn.init_async()  # type: ignore[attr-defined]
+                connectors[name] = conn
+
+            services = {name: cfg.service(*cfg.args, **cfg.kwargs) for name, cfg in config.services.items()}
+
+            instance = config.task(
+                connectors=connectors,
+                services=services,
+                policies=config.policies,
+                worker_id=worker_id,
+                task_config=config,
+                *args,
+                **kwargs,
+            )
+
+            if mode == "transaction":
+                await __run_transaction_async(
+                    instance=instance,
+                    policy=kwargs.get("policy"),  # type: ignore[arg-type]
+                    transaction=kwargs.get("transaction"),
+                    transaction_id=kwargs.get("transaction_id"),
+                )
+            else:
+                await instance.execute()
+
+    finally:
+        if instance is not None:
+            await instance.exit()
+
+
+# =============================================================
+# SYNC TRANSACTION EXECUTION
+# =============================================================
+
+
+def __run_transaction_sync(
+    instance: BaseTask,
+    policy: str,
+    transaction: Transaction | None = None,
+    transaction_id: str | None = None,
+):
+    policy_obj = next((p for p in instance.policies if p.name == policy), None)
+    if not policy_obj:
+        raise ValueError(f"Policy '{policy}' not found")
+
+    if not transaction and not transaction_id:
+        raise ValueError("Either transaction or transaction_id must be provided")
+
+    if transaction_id:
+        conn = instance._resolve_connector(policy_obj.fetch.connector_name)  # type: ignore[attr-defined]
+        transaction = conn.fetch_transaction_by_id(transaction_id)
+
+    success, result = instance._start_processing(transaction, policy_obj)  # type: ignore[attr-defined]
+    if not success:
+        raise result
+    return result
+
+
+# =============================================================
+# SYNC TASK EXECUTION
+# =============================================================
+
+
+def __run_task_sync(
+    config: TaskConfig,
+    mode: Literal["task", "transaction"] = "task",
+    *args,
+    **kwargs,
+):
+    if not issubclass(config.task, BaseTask):  # type: ignore[arg-type]
+        raise ValueError(f"Invalid sync task: {config.task}")
+
+    worker_id = kwargs.pop("worker_id", None)
+
+    execution_start_time = datetime.now().isoformat()
+    task_exec_metadata = TaskExecMetadata(
+        task_name=config.name,
+        execution_id=str(uuid.uuid4()),
+        execution_start_time=execution_start_time,
+        pid=os.getpid(),
+        worker_id=worker_id,
+    ).model_dump()
+
+    __init_telemetry(config, task=config.task, task_exec_metadata=task_exec_metadata)
+    tracer = tracing.get_tracer(__name__)
+    logger = logging.get_logger(__name__)
+
+    instance = None
+
+    logger.info(f"Task {config.name} started at {execution_start_time}.")
+
+    with tracer.start_as_current_span(
+        f"{config.task.__name__}.run",
+        attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+    ):
+        try:
+            logger.info("Initializing connectors...")
+            connectors = {}
+            with tracer.start_as_current_span(
+                f"{config.task.__name__}.connectors.init",
+                attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+            ):
+                for name, cfg in config.connectors.items():
+                    with tracer.start_as_current_span(
+                        f"{config.task.__name__}.connectors.{name}.init",
+                        attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+                    ):
+                        connectors[name] = cfg.connector(*cfg.args, **cfg.kwargs)
+
+            logger.info("Initializing services...")
+            services = {}
+            with tracer.start_as_current_span(
+                f"{config.task.__name__}.services.init",
+                attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+            ):
+                for name, cfg in config.services.items():
+                    with tracer.start_as_current_span(
+                        f"{config.task.__name__}.services.{name}.init",
+                        attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+                    ):
+                        services[name] = cfg.service(*cfg.args, **cfg.kwargs)
+
+            with tracer.start_as_current_span(
+                f"{config.task.__name__}.instance.init",
+                attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+            ):
+                logger.info("Creating task instance...")
+                instance = config.task(
+                    connectors=connectors,
+                    services=services,
+                    policies=config.policies,
+                    worker_id=worker_id,
+                    task_config=config,
+                    *args,
+                    **kwargs,
+                )
+
+            if mode == "transaction":
+                logger.info("Running task in transaction execution mode...")
+                __run_transaction_sync(
+                    instance=instance,
+                    policy=kwargs.get("policy"),  # type: ignore[arg-type]
+                    transaction=kwargs.get("transaction", None),
+                    transaction_id=kwargs.get("transaction_id", None),
+                )
+            else:
+                with tracer.start_as_current_span(
+                    f"{config.task.__name__}.execute",
+                    attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+                ):
+                    logger.info("Running task in full execution mode...")
+                    instance.execute()
+
+        finally:
+            if instance is not None:
+                with tracer.start_as_current_span(
+                    f"{config.task.__name__}.exit",
+                    attributes={"task.execution.id": task_exec_metadata["execution_id"]},
+                ):
+                    logger.info(f"Exiting task {config.name}...")
+                    instance.exit()
+
+
+# =============================================================
+# PUBLIC API — RUNNER
+# =============================================================
+
+
+def run_task(
+    config: TaskConfig,
+    debug: bool = False,
+    mode: Literal["task", "transaction"] = "task",
+    *args,
+    **kwargs,
+) -> int:
+    """
+    Process entrypoint.
+
+    Returns POSIX-compatible exit code.
+    """
+
+    _install_signal_handlers()
+    is_async = issubclass(config.task, BaseAsyncTask)  # type: ignore[arg-type]
+
+    # ---------------------------------------------------------
+    # SINGLE PROCESS
+    # ---------------------------------------------------------
+    if debug or config.max_workers == 1:
+        try:
+            if is_async:
+                asyncio.run(__run_task_async(config, mode, *args, **kwargs))
+            else:
+                __run_task_sync(config, mode, *args, **kwargs)
+
+            if is_shutdown_requested():
+                return int(get_shutdown_exit_code())
+
+            return int(ExitCode.SUCCESS)
+
+        except ValueError:
+            traceback.print_exc()
+            return int(ExitCode.CONFIG_ERROR)
+
+        except Exception:
+            traceback.print_exc()
+            return int(ExitCode.ERROR)
+
+    # ---------------------------------------------------------
+    # MULTI-PROCESS (SYNC ONLY)
+    # ---------------------------------------------------------
+    if is_async:
+        raise RuntimeError("Async tasks cannot be executed with multiple processes. Use debug=True or max_workers=1.")
+
+    has_error = False
+
+    with ProcessPoolExecutor(max_workers=config.max_workers) as executor:
+        futures = []
+        for worker_id in range(config.max_workers):
+            worker_kwargs = {
+                **kwargs,
+                "worker_id": worker_id,
+                "total_workers": config.max_workers,
+            }
+            futures.append(executor.submit(__run_task_sync, config, mode, *args, **worker_kwargs))
+
+        for f in futures:
+            try:
+                f.result()
+            except Exception:
+                traceback.print_exc()
+                has_error = True
+
+    if is_shutdown_requested():
+        return int(get_shutdown_exit_code())
+
+    if has_error:
+        return int(ExitCode.ERROR)
+
+    return int(ExitCode.SUCCESS)

ergon/task/utils.py

@@ -0,0 +1,64 @@
+# utils.py
+import asyncio
+import logging
+import math
+import time
+from datetime import datetime
+
+from .. import telemetry
+
+logger = logging.getLogger(__name__)
+tracer = telemetry.tracing.get_tracer(__name__)
+
+
+def _get_wake_time_iso(delay: float) -> str:
+    return datetime.fromtimestamp(time.time() + delay).isoformat()
+
+
+# ============================================================
+#  BACKOFF / SLEEP HELPERS
+# ============================================================
+def compute_backoff(backoff: float, multiplier: float, cap: float, attempt: int) -> float:
+    with tracer.start_as_current_span("compute_backoff"):
+        logger.info(
+            f"Computing backoff with arguments: "
+            f"attempt {attempt}, "
+            f"backoff {backoff}, "
+            f"multiplier {multiplier}, "
+            f"and cap {cap}"
+        )
+
+        if cap > 0 and multiplier > 1 and backoff > 0:
+            # max attempt that won't exceed cap
+            max_attempt = math.floor(math.log(cap / backoff, multiplier)) if cap > backoff else 0
+            safe_attempt = min(attempt, max_attempt)
+        else:
+            safe_attempt = attempt
+
+        delay = backoff * (multiplier**safe_attempt)
+        computed_delay = min(delay, cap) if cap > 0 else delay
+
+        logger.info(f"Computed backoff: {computed_delay} seconds")
+        return computed_delay
+
+
+def backoff(backoff: float, multiplier: float, cap: float, attempt: int):
+    """Blocking sleep with computed backoff."""
+    delay = compute_backoff(backoff, multiplier, cap, attempt)
+    if delay > 0:
+        with tracer.start_as_current_span("sleep", attributes={"delay": delay}):
+            estimated_wake_time_iso = _get_wake_time_iso(delay)
+            logger.info(f"Sleeping for {delay} seconds until {estimated_wake_time_iso}")
+            time.sleep(delay)
+            logger.info(f"Woke up from {delay} second{'' if delay == 1 else 's'} sleep")
+
+
+async def backoff_async(backoff: float, multiplier: float, cap: float, attempt: int):
+    """Async backoff with computed backoff."""
+    delay = compute_backoff(backoff, multiplier, cap, attempt)
+    if delay > 0:
+        with tracer.start_as_current_span("sleep", attributes={"delay": delay}):
+            estimated_wake_time_iso = _get_wake_time_iso(delay)
+            logger.info(f"Sleeping for {delay} seconds until {estimated_wake_time_iso}")
+            await asyncio.sleep(delay)
+            logger.info(f"Woke up from {delay} second{'' if delay == 1 else 's'} sleep")

ergon/telemetry/__init__.py

@@ -0,0 +1,7 @@
+from . import logging, metrics, tracing
+
+__all__ = [
+    "logging",
+    "metrics",
+    "tracing",
+]

ergon/telemetry/_resource.py

@@ -0,0 +1,13 @@
+def _inject_otel_resource_attributes(resource: dict, metadata: dict) -> dict:
+    enriched = dict(resource)
+    enriched.update(
+        {
+            "ergon.task.name": metadata["task_name"],
+            "ergon.task.execution.id": metadata["execution_id"],
+            "ergon.task.execution.pid": metadata["pid"],
+            "ergon.task.execution.host.name": metadata["host_name"],
+            "ergon.task.execution.host.ip": metadata["host_ip"],
+            "ergon.task.execution.start_time": metadata["execution_start_time"],
+        }
+    )
+    return enriched