avtomatika 1.0b6__py3-none-any.whl → 1.0b8__py3-none-any.whl

avtomatika/executor.py

@@ -1,5 +1,4 @@
 from asyncio import CancelledError, Task, create_task, sleep
-from inspect import signature
 from logging import getLogger
 from time import monotonic
 from types import SimpleNamespace
@@ -48,6 +47,7 @@ except ImportError:
     inject = NoOpPropagate().inject
     TraceContextTextMapPropagator = NoOpTraceContextTextMapPropagator()  # Instantiate the class
 
+from .app_keys import S3_SERVICE_KEY
 from .context import ActionFactory
 from .data_types import ClientConfig, JobContext
 from .history.base import HistoryStorageBase
@@ -75,7 +75,7 @@ class JobExecutor:
         self._running = False
         self._processing_messages: set[str] = set()
 
-    async def _process_job(self, job_id: str, message_id: str):
+    async def _process_job(self, job_id: str, message_id: str) -> None:
         """The core logic for processing a single job dequeued from storage."""
         if message_id in self._processing_messages:
             return
@@ -144,6 +144,11 @@ class JobExecutor:
                     plan=client_config_dict.get("plan", "unknown"),
                     params=client_config_dict.get("params", {}),
                 )
+
+                # Get TaskFiles if S3 service is available
+                s3_service = self.engine.app.get(S3_SERVICE_KEY)
+                task_files = s3_service.get_task_files(job_id) if s3_service else None
+
                 context = JobContext(
                     job_id=job_id,
                     current_state=job_state["current_state"],
@@ -154,6 +159,7 @@ class JobExecutor:
                     data_stores=SimpleNamespace(**blueprint.data_stores),
                     tracing_context=tracing_context,
                     aggregation_results=job_state.get("aggregation_results"),
+                    task_files=task_files,
                 )
 
                 try:
@@ -167,19 +173,24 @@ class JobExecutor:
                         handler = blueprint.find_handler(context.current_state, context)
 
                     # Build arguments for the handler dynamically.
-                    handler_signature = signature(handler)
-                    params_to_inject = {}
+                    param_names = blueprint.get_handler_params(handler)
+                    params_to_inject: dict[str, Any] = {}
 
-                    if "context" in handler_signature.parameters:
+                    if "context" in param_names:
                         params_to_inject["context"] = context
-                        if "actions" in handler_signature.parameters:
+                        if "actions" in param_names:
                             params_to_inject["actions"] = action_factory
+                        if "task_files" in param_names:
+                            params_to_inject["task_files"] = task_files
                     else:
                         # New injection logic with prioritized lookup.
                         context_as_dict = context._asdict()
-                        for param_name in handler_signature.parameters:
+                        for param_name in param_names:
+                            # Direct injection of task_files
+                            if param_name == "task_files":
+                                params_to_inject[param_name] = task_files
                             # Look in JobContext fields first.
-                            if param_name in context_as_dict:
+                            elif param_name in context_as_dict:
                                 params_to_inject[param_name] = context_as_dict[param_name]
                             # Then look in state_history (data from previous steps/workers).
                             elif param_name in context.state_history:
@@ -232,7 +243,7 @@ class JobExecutor:
         job_state: dict[str, Any],
         next_state: str,
         duration_ms: int,
-    ):
+    ) -> None:
         job_id = job_state["id"]
         previous_state = job_state["current_state"]
         logger.info(f"Job {job_id} transitioning from {previous_state} to {next_state}")
@@ -259,14 +270,29 @@ class JobExecutor:
             await self.storage.enqueue_job(job_id)
         else:
             logger.info(f"Job {job_id} reached terminal state {next_state}")
+
+            # Clean up S3 files if service is available
+            s3_service = self.engine.app.get(S3_SERVICE_KEY)
+            if s3_service:
+                task_files = s3_service.get_task_files(job_id)
+                if task_files:
+                    # Run cleanup in background to not block response
+                    create_task(task_files.cleanup())
+
             await self._check_and_resume_parent(job_state)
+            # Send webhook for finished/failed jobs
+            event_type = "job_finished" if next_state == "finished" else "job_failed"
+            # Since _check_and_resume_parent is for sub-jobs, we only send webhook if it's a top-level job
+            # or if the user explicitly requested it for sub-jobs (by providing webhook_url).
+            # The current logic stores webhook_url in job_state, so we just check it.
+            await self.engine.send_job_webhook(job_state, event_type)
 
     async def _handle_dispatch(
         self,
         job_state: dict[str, Any],
         task_info: dict[str, Any],
         duration_ms: int,
-    ):
+    ) -> None:
         job_id = job_state["id"]
         current_state = job_state["current_state"]
 
@@ -302,7 +328,6 @@ class JobExecutor:
             await self.storage.save_job_state(job_id, job_state)
             await self.storage.add_job_to_watch(job_id, timeout_at)
 
-            # Now, dispatch the task
             await self.dispatcher.dispatch(job_state, task_info)
 
     async def _handle_run_blueprint(
@@ -310,7 +335,7 @@ class JobExecutor:
         parent_job_state: dict[str, Any],
         sub_blueprint_info: dict[str, Any],
         duration_ms: int,
-    ):
+    ) -> None:
         parent_job_id = parent_job_state["id"]
         child_job_id = str(uuid4())
 
@@ -350,7 +375,7 @@ class JobExecutor:
         job_state: dict[str, Any],
         parallel_info: dict[str, Any],
         duration_ms: int,
-    ):
+    ) -> None:
         job_id = job_state["id"]
         tasks_to_dispatch = parallel_info["tasks"]
         aggregate_into = parallel_info["aggregate_into"]
@@ -398,7 +423,7 @@ class JobExecutor:
         job_state: dict[str, Any],
         error: Exception,
         duration_ms: int,
-    ):
+    ) -> None:
         """Handles failures that occur *during the execution of a handler*.
 
         This is different from a task failure reported by a worker. This logic
@@ -447,13 +472,14 @@ class JobExecutor:
             await self.storage.quarantine_job(job_id)
             # If this quarantined job was a sub-job, we must now resume its parent.
             await self._check_and_resume_parent(job_state)
+            await self.engine.send_job_webhook(job_state, "job_quarantined")
             from . import metrics
 
             metrics.jobs_failed_total.inc(
                 {metrics.LABEL_BLUEPRINT: job_state.get("blueprint_name", "unknown")},
             )
 
-    async def _check_and_resume_parent(self, child_job_state: dict[str, Any]):
+    async def _check_and_resume_parent(self, child_job_state: dict[str, Any]) -> None:
         """Checks if a completed job was a sub-job. If so, it resumes the parent
         job, passing the success/failure outcome of the child.
         """
@@ -493,7 +519,7 @@ class JobExecutor:
         await self.storage.enqueue_job(parent_job_id)
 
     @staticmethod
-    def _handle_task_completion(task: Task):
+    def _handle_task_completion(task: Task) -> None:
         """Callback to handle completion of a job processing task."""
         try:
             # This will re-raise any exception caught in the task
@@ -505,7 +531,7 @@ class JobExecutor:
             # Log any other exceptions that occurred in the task.
             logger.exception("Unhandled exception in job processing task")
 
-    async def run(self):
+    async def run(self) -> None:
         import asyncio
 
         logger.info("JobExecutor started.")
@@ -517,7 +543,10 @@ class JobExecutor:
                 # Wait for an available slot before fetching a new job
                 await semaphore.acquire()
 
-                result = await self.storage.dequeue_job()
+                # Block for a configured time waiting for a job
+                block_time = self.engine.config.REDIS_STREAM_BLOCK_MS
+                result = await self.storage.dequeue_job(block=block_time if block_time > 0 else None)
+
                 if result:
                     job_id, message_id = result
                     task = create_task(self._process_job(job_id, message_id))
@@ -525,16 +554,20 @@ class JobExecutor:
                     # Release the semaphore slot when the task is done
                     task.add_done_callback(lambda _: semaphore.release())
                 else:
-                    # No job found, release the slot and wait a bit
+                    # Timeout reached, release slot and loop again
                     semaphore.release()
-                    # Prevent busy loop if storage returns None immediately
-                    await sleep(0.1)
+                    # Prevent busy loop if blocking is disabled (e.g. in tests) or failed
+                    if block_time <= 0:
+                        await sleep(0.1)
+
             except CancelledError:
                 break
             except Exception:
                 logger.exception("Error in JobExecutor main loop.")
+                # If an error occurred (e.g. Redis connection lost), sleep briefly to avoid log spam
+                semaphore.release()
                 await sleep(1)
         logger.info("JobExecutor stopped.")
 
-    def stop(self):
+    def stop(self) -> None:
         self._running = False

avtomatika/health_checker.py

@@ -20,19 +20,37 @@ logger = getLogger(__name__)
 
 
 class HealthChecker:
-    def __init__(self, engine: "OrchestratorEngine"):
+    def __init__(self, engine: "OrchestratorEngine", interval_seconds: int = 600):
+        self.engine = engine
+        self.storage = engine.storage
+        self.interval_seconds = interval_seconds
         self._running = False
+        from uuid import uuid4
+
+        self._instance_id = str(uuid4())
 
     async def run(self):
-        logger.info("HealthChecker is now passive and will not perform active checks.")
+        logger.info(f"HealthChecker started (Active Index Cleanup, Instance ID: {self._instance_id}).")
         self._running = True
         while self._running:
             try:
-                # Sleep for a long time, as this checker is passive.
-                # The loop exists to allow for a clean shutdown.
-                await sleep(3600)
+                # Use distributed lock to ensure only one instance cleans up
+                if await self.storage.acquire_lock(
+                    "global_health_check_lock", self._instance_id, self.interval_seconds - 5
+                ):
+                    try:
+                        await self.storage.cleanup_expired_workers()
+                    finally:
+                        # We don't release the lock immediately to prevent other instances from
+                        # running the same task if the interval is small.
+                        pass
+
+                await sleep(self.interval_seconds)
             except CancelledError:
                 break
+            except Exception:
+                logger.exception("Error in HealthChecker main loop.")
+                await sleep(60)
         logger.info("HealthChecker stopped.")
 
     def stop(self):

avtomatika/history/base.py

@@ -1,25 +1,79 @@
+import asyncio
+import contextlib
 from abc import ABC, abstractmethod
+from logging import getLogger
 from typing import Any
 
+logger = getLogger(__name__)
+
 
 class HistoryStorageBase(ABC):
     """Abstract base class for a history store.
-    Defines the interface for logging job and worker events.
+    Implements buffered asynchronous logging to avoid blocking the main loop.
     """
 
+    def __init__(self):
+        self._queue: asyncio.Queue[tuple[str, dict[str, Any]]] = asyncio.Queue(maxsize=5000)
+        self._worker_task: asyncio.Task | None = None
+
+    async def start(self) -> None:
+        """Starts the background worker for writing logs."""
+        if not self._worker_task:
+            self._worker_task = asyncio.create_task(self._worker())
+            logger.info("HistoryStorage background worker started.")
+
+    async def close(self) -> None:
+        """Stops the background worker and closes resources."""
+        if self._worker_task:
+            self._worker_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._worker_task
+            self._worker_task = None
+            logger.info("HistoryStorage background worker stopped.")
+
     @abstractmethod
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Performs initialization, e.g., creating tables in the DB."""
         raise NotImplementedError
 
+    async def log_job_event(self, event_data: dict[str, Any]) -> None:
+        """Queues a job event for logging."""
+        try:
+            self._queue.put_nowait(("job", event_data))
+        except asyncio.QueueFull:
+            logger.warning("History queue full! Dropping job event.")
+
+    async def log_worker_event(self, event_data: dict[str, Any]) -> None:
+        """Queues a worker event for logging."""
+        try:
+            self._queue.put_nowait(("worker", event_data))
+        except asyncio.QueueFull:
+            logger.warning("History queue full! Dropping worker event.")
+
+    async def _worker(self) -> None:
+        while True:
+            try:
+                kind, data = await self._queue.get()
+                try:
+                    if kind == "job":
+                        await self._persist_job_event(data)
+                    elif kind == "worker":
+                        await self._persist_worker_event(data)
+                except Exception as e:
+                    logger.error(f"Error persisting history event: {e}")
+                finally:
+                    self._queue.task_done()
+            except asyncio.CancelledError:
+                break
+
     @abstractmethod
-    async def log_job_event(self, event_data: dict[str, Any]):
-        """Logs an event related to the job lifecycle."""
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
+        """Actual implementation of writing a job event to storage."""
         raise NotImplementedError
 
     @abstractmethod
-    async def log_worker_event(self, event_data: dict[str, Any]):
-        """Logs an event related to the worker lifecycle."""
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
+        """Actual implementation of writing a worker event to storage."""
         raise NotImplementedError
 
     @abstractmethod

avtomatika/history/noop.py

@@ -8,20 +8,31 @@ class NoOpHistoryStorage(HistoryStorageBase):
     Used when history storage is not configured.
     """
 
-    async def initialize(self):
-        # Do nothing
+    def __init__(self):
+        super().__init__()
+
+    async def start(self) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass
+
+    async def initialize(self) -> None:
+        pass
+
+    async def log_job_event(self, event_data: dict[str, Any]) -> None:
+        pass
+
+    async def log_worker_event(self, event_data: dict[str, Any]) -> None:
         pass
 
-    async def log_job_event(self, event_data: dict[str, Any]):
-        # Do nothing
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
         pass
 
-    async def log_worker_event(self, event_data: dict[str, Any]):
-        # Do nothing
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
         pass
 
     async def get_job_history(self, job_id: str) -> list[dict[str, Any]]:
-        # Always return an empty list
         return []
 
     async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:

avtomatika/history/postgres.py

@@ -46,19 +46,20 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
     """Implementation of the history store based on asyncpg for PostgreSQL."""
 
     def __init__(self, dsn: str, tz_name: str = "UTC"):
+        super().__init__()
         self._dsn = dsn
         self._pool: Pool | None = None
         self.tz_name = tz_name
         self.tz = ZoneInfo(tz_name)
 
-    async def _setup_connection(self, conn: Connection):
+    async def _setup_connection(self, conn: Connection) -> None:
         """Configures the connection session with the correct timezone."""
         try:
             await conn.execute(f"SET TIME ZONE '{self.tz_name}'")
         except PostgresError as e:
             logger.error(f"Failed to set timezone '{self.tz_name}' for PG connection: {e}")
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Initializes the connection pool to PostgreSQL and creates tables."""
         try:
             # We use init parameter to configure each new connection in the pool
@@ -75,13 +76,14 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
             logger.error(f"Failed to initialize PostgreSQL history storage: {e}")
             raise
 
-    async def close(self):
-        """Closes the connection pool."""
+    async def close(self) -> None:
+        """Closes the connection pool and background worker."""
+        await super().close()
         if self._pool:
             await self._pool.close()
             logger.info("PostgreSQL history storage connection pool closed.")
 
-    async def log_job_event(self, event_data: dict[str, Any]):
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
         """Logs a job lifecycle event to PostgreSQL."""
         if not self._pool:
             raise RuntimeError("History storage is not initialized.")
@@ -117,7 +119,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
         except PostgresError as e:
             logger.error(f"Failed to log job event to PostgreSQL: {e}")
 
-    async def log_worker_event(self, event_data: dict[str, Any]):
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
         """Logs a worker lifecycle event to PostgreSQL."""
         if not self._pool:
             raise RuntimeError("History storage is not initialized.")

avtomatika/history/sqlite.py

@@ -49,11 +49,12 @@ class SQLiteHistoryStorage(HistoryStorageBase):
     """
 
     def __init__(self, db_path: str, tz_name: str = "UTC"):
+        super().__init__()
         self._db_path = db_path
         self._conn: Connection | None = None
         self.tz = ZoneInfo(tz_name)
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Initializes the database connection and creates tables if they don't exist."""
         try:
             self._conn = await connect(self._db_path)
@@ -68,8 +69,9 @@ class SQLiteHistoryStorage(HistoryStorageBase):
             logger.error(f"Failed to initialize SQLite history storage: {e}")
             raise
 
-    async def close(self):
-        """Closes the database connection."""
+    async def close(self) -> None:
+        """Closes the database connection and background worker."""
+        await super().close()
         if self._conn:
             await self._conn.close()
             logger.info("SQLite history storage connection closed.")
@@ -91,7 +93,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
 
         return item
 
-    async def log_job_event(self, event_data: dict[str, Any]):
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
         """Logs a job lifecycle event to the job_history table."""
         if not self._conn:
             raise RuntimeError("History storage is not initialized.")
@@ -128,7 +130,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
         except Error as e:
             logger.error(f"Failed to log job event: {e}")
 
-    async def log_worker_event(self, event_data: dict[str, Any]):
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
         """Logs a worker lifecycle event to the worker_history table."""
         if not self._conn:
             raise RuntimeError("History storage is not initialized.")

avtomatika/metrics.py

@@ -12,7 +12,7 @@ task_queue_length: Gauge
 active_workers: Gauge
 
 
-def init_metrics():
+def init_metrics() -> None:
     """
     Initializes Prometheus metrics.
     Uses a registry check for idempotency, which is important for tests.

avtomatika/reputation.py

@@ -52,48 +52,54 @@ class ReputationCalculator:
     async def calculate_all_reputations(self):
         """Calculates and updates the reputation for all active workers."""
         logger.info("Starting reputation calculation for all workers...")
-        workers = await self.storage.get_available_workers()
-        if not workers:
+
+        # Get only IDs of active workers to avoid O(N) scan of all data
+        worker_ids = await self.storage.get_active_worker_ids()
+
+        if not worker_ids:
             logger.info("No active workers found for reputation calculation.")
             return
 
-        for worker in workers:
-            worker_id = worker.get("worker_id")
-            if not worker_id:
-                continue
-
-            history = await self.history_storage.get_worker_history(
-                worker_id,
-                since_days=REPUTATION_HISTORY_DAYS,
-            )
-
-            # Count only task completion events
-            task_finished_events = [event for event in history if event.get("event_type") == "task_finished"]
-
-            if not task_finished_events:
-                # If there is no history, the reputation does not change (remains 1.0 by default)
-                continue
-
-            successful_tasks = 0
-            for event in task_finished_events:
-                # Extract the result from the snapshot
-                snapshot = event.get("context_snapshot", {})
-                result = snapshot.get("result", {})
-                if result.get("status") == "success":
-                    successful_tasks += 1
-
-            total_tasks = len(task_finished_events)
-            new_reputation = successful_tasks / total_tasks if total_tasks > 0 else 1.0
-
-            # Round for cleanliness
-            new_reputation = round(new_reputation, 4)
-
-            logger.info(
-                f"Updating reputation for worker {worker_id}: {worker.get('reputation')} -> {new_reputation}",
-            )
-            await self.storage.update_worker_data(
-                worker_id,
-                {"reputation": new_reputation},
-            )
+        logger.info(f"Recalculating reputation for {len(worker_ids)} workers.")
+
+        for worker_id in worker_ids:
+            if not self._running:
+                break
+
+            try:
+                history = await self.history_storage.get_worker_history(
+                    worker_id,
+                    since_days=REPUTATION_HISTORY_DAYS,
+                )
+
+                # Count only task completion events
+                task_finished_events = [event for event in history if event.get("event_type") == "task_finished"]
+
+                if not task_finished_events:
+                    # If there is no history, skip to next worker
+                    continue
+
+                successful_tasks = 0
+                for event in task_finished_events:
+                    # Extract the result from the snapshot
+                    snapshot = event.get("context_snapshot", {})
+                    result = snapshot.get("result", {})
+                    if result.get("status") == "success":
+                        successful_tasks += 1
+
+                total_tasks = len(task_finished_events)
+                new_reputation = successful_tasks / total_tasks if total_tasks > 0 else 1.0
+                new_reputation = round(new_reputation, 4)
+
+                await self.storage.update_worker_data(
+                    worker_id,
+                    {"reputation": new_reputation},
+                )
+
+                # Throttling: Small sleep to prevent DB spikes
+                await sleep(0.1)
+
+            except Exception as e:
+                logger.error(f"Failed to calculate reputation for worker {worker_id}: {e}")
 
         logger.info("Reputation calculation finished.")