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

avtomatika/executor.py

@@ -3,7 +3,7 @@ from inspect import signature
 from logging import getLogger
 from time import monotonic
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 # Conditional import for OpenTelemetry
@@ -73,158 +73,163 @@ class JobExecutor:
         self.history_storage = history_storage
         self.dispatcher = engine.dispatcher
         self._running = False
+        self._processing_messages: set[str] = set()
 
-    async def _process_job(self, job_id: str):
-        """The core logic for processing a single job dequeued from storage.
-
-        This function orchestrates finding the correct blueprint and handler,
-        executing the handler, and then processing the action (e.g., transition,
-        dispatch) that the handler requested.
-        """
-        start_time = monotonic()
-        job_state = await self.storage.get_job_state(job_id)
-        if not job_state:
-            logger.error(f"Job {job_id} not found in storage, cannot process.")
+    async def _process_job(self, job_id: str, message_id: str):
+        """The core logic for processing a single job dequeued from storage."""
+        if message_id in self._processing_messages:
             return
 
-        if job_state.get("status") in TERMINAL_STATES:
-            logger.warning(f"Job {job_id} is already in a terminal state '{job_state['status']}', skipping.")
-            return
-
-        # Ensure retry_count is initialized.
-        if "retry_count" not in job_state:
-            job_state["retry_count"] = 0
-
-        await self.history_storage.log_job_event(
-            {
-                "job_id": job_id,
-                "state": job_state.get("current_state"),
-                "event_type": "state_started",
-                "attempt_number": job_state.get("retry_count", 0) + 1,
-                "context_snapshot": job_state,
-            },
-        )
-
-        # Set up distributed tracing context.
-        parent_context = TraceContextTextMapPropagator().extract(
-            carrier=job_state.get("tracing_context", {}),
-        )
+        self._processing_messages.add(message_id)
+        try:
+            start_time = monotonic()
+            job_state = await self.storage.get_job_state(job_id)
+            if not job_state:
+                logger.error(f"Job {job_id} not found in storage, cannot process.")
+                return
 
-        with tracer.start_as_current_span(
-            f"JobExecutor:{job_state['blueprint_name']}:{job_state['current_state']}",
-            context=parent_context,
-        ) as span:
-            span.set_attribute("job.id", job_id)
-            span.set_attribute("job.current_state", job_state["current_state"])
-
-            # Inject the current tracing context back into the job state for propagation.
-            tracing_context: Dict[str, str] = {}
-            inject(tracing_context)
-            job_state["tracing_context"] = tracing_context
-
-            blueprint = self.engine.blueprints.get(job_state["blueprint_name"])
-            if not blueprint:
-                # This is a critical, non-retriable error.
-                duration_ms = int((monotonic() - start_time) * 1000)
-                await self._handle_failure(
-                    job_state,
-                    RuntimeError(
-                        f"Blueprint '{job_state['blueprint_name']}' not found",
-                    ),
-                    duration_ms,
-                )
+            if job_state.get("status") in TERMINAL_STATES:
+                logger.warning(f"Job {job_id} is already in a terminal state '{job_state['status']}', skipping.")
                 return
 
-            # Prepare the context and action factory for the handler.
-            action_factory = ActionFactory(job_id)
-            client_config_dict = job_state.get("client_config", {})
-            client_config = ClientConfig(
-                token=client_config_dict.get("token", ""),
-                plan=client_config_dict.get("plan", "unknown"),
-                params=client_config_dict.get("params", {}),
-            )
-            context = JobContext(
-                job_id=job_id,
-                current_state=job_state["current_state"],
-                initial_data=job_state["initial_data"],
-                state_history=job_state.get("state_history", {}),
-                client=client_config,
-                actions=action_factory,
-                data_stores=SimpleNamespace(**blueprint.data_stores),
-                tracing_context=tracing_context,
-                aggregation_results=job_state.get("aggregation_results"),
+            # Ensure retry_count is initialized.
+            if "retry_count" not in job_state:
+                job_state["retry_count"] = 0
+
+            await self.history_storage.log_job_event(
+                {
+                    "job_id": job_id,
+                    "state": job_state.get("current_state"),
+                    "event_type": "state_started",
+                    "attempt_number": job_state.get("retry_count", 0) + 1,
+                    "context_snapshot": job_state,
+                },
             )
 
-            try:
-                # Find and execute the appropriate handler for the current state.
-                # It's important to check for aggregator handlers first for states
-                # that are targets of parallel execution.
-                is_aggregator_state = job_state.get("aggregation_target") == job_state.get("current_state")
-                if is_aggregator_state and job_state.get("current_state") in blueprint.aggregator_handlers:
-                    handler = blueprint.aggregator_handlers[job_state["current_state"]]
-                else:
-                    handler = blueprint.find_handler(context.current_state, context)
-
-                # Build arguments for the handler dynamically.
-                handler_signature = signature(handler)
-                params_to_inject = {}
+            # Set up distributed tracing context.
+            parent_context = TraceContextTextMapPropagator().extract(
+                carrier=job_state.get("tracing_context", {}),
+            )
 
-                if "context" in handler_signature.parameters:
-                    params_to_inject["context"] = context
-                    if "actions" in handler_signature.parameters:
-                        params_to_inject["actions"] = action_factory
-                else:
-                    # New injection logic with prioritized lookup.
-                    context_as_dict = context._asdict()
-                    for param_name in handler_signature.parameters:
-                        # Look in JobContext fields first.
-                        if 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:
-                            params_to_inject[param_name] = context.state_history[param_name]
-                        # Finally, look in the initial data the job was created with.
-                        elif param_name in context.initial_data:
-                            params_to_inject[param_name] = context.initial_data[param_name]
-
-                await handler(**params_to_inject)
-
-                duration_ms = int((monotonic() - start_time) * 1000)
-
-                # Process the single action requested by the handler.
-                if action_factory.next_state:
-                    await self._handle_transition(
-                        job_state,
-                        action_factory.next_state,
-                        duration_ms,
-                    )
-                elif action_factory.task_to_dispatch:
-                    await self._handle_dispatch(
-                        job_state,
-                        action_factory.task_to_dispatch,
-                        duration_ms,
-                    )
-                elif action_factory.parallel_tasks_to_dispatch:
-                    await self._handle_parallel_dispatch(
+            with tracer.start_as_current_span(
+                f"JobExecutor:{job_state['blueprint_name']}:{job_state['current_state']}",
+                context=parent_context,
+            ) as span:
+                span.set_attribute("job.id", job_id)
+                span.set_attribute("job.current_state", job_state["current_state"])
+
+                # Inject the current tracing context back into the job state for propagation.
+                tracing_context: dict[str, str] = {}
+                inject(tracing_context)
+                job_state["tracing_context"] = tracing_context
+
+                blueprint = self.engine.blueprints.get(job_state["blueprint_name"])
+                if not blueprint:
+                    # This is a critical, non-retriable error.
+                    duration_ms = int((monotonic() - start_time) * 1000)
+                    await self._handle_failure(
                         job_state,
-                        action_factory.parallel_tasks_to_dispatch,
-                        duration_ms,
-                    )
-                elif action_factory.sub_blueprint_to_run:
-                    await self._handle_run_blueprint(
-                        job_state,
-                        action_factory.sub_blueprint_to_run,
+                        RuntimeError(
+                            f"Blueprint '{job_state['blueprint_name']}' not found",
+                        ),
                         duration_ms,
                     )
+                    return
+
+                # Prepare the context and action factory for the handler.
+                action_factory = ActionFactory(job_id)
+                client_config_dict = job_state.get("client_config", {})
+                client_config = ClientConfig(
+                    token=client_config_dict.get("token", ""),
+                    plan=client_config_dict.get("plan", "unknown"),
+                    params=client_config_dict.get("params", {}),
+                )
+                context = JobContext(
+                    job_id=job_id,
+                    current_state=job_state["current_state"],
+                    initial_data=job_state["initial_data"],
+                    state_history=job_state.get("state_history", {}),
+                    client=client_config,
+                    actions=action_factory,
+                    data_stores=SimpleNamespace(**blueprint.data_stores),
+                    tracing_context=tracing_context,
+                    aggregation_results=job_state.get("aggregation_results"),
+                )
 
-            except Exception as e:
-                # This catches errors within the handler's execution.
-                duration_ms = int((monotonic() - start_time) * 1000)
-                await self._handle_failure(job_state, e, duration_ms)
+                try:
+                    # Find and execute the appropriate handler for the current state.
+                    # It's important to check for aggregator handlers first for states
+                    # that are targets of parallel execution.
+                    is_aggregator_state = job_state.get("aggregation_target") == job_state.get("current_state")
+                    if is_aggregator_state and job_state.get("current_state") in blueprint.aggregator_handlers:
+                        handler = blueprint.aggregator_handlers[job_state["current_state"]]
+                    else:
+                        handler = blueprint.find_handler(context.current_state, context)
+
+                    # Build arguments for the handler dynamically.
+                    handler_signature = signature(handler)
+                    params_to_inject = {}
+
+                    if "context" in handler_signature.parameters:
+                        params_to_inject["context"] = context
+                        if "actions" in handler_signature.parameters:
+                            params_to_inject["actions"] = action_factory
+                    else:
+                        # New injection logic with prioritized lookup.
+                        context_as_dict = context._asdict()
+                        for param_name in handler_signature.parameters:
+                            # Look in JobContext fields first.
+                            if 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:
+                                params_to_inject[param_name] = context.state_history[param_name]
+                            # Finally, look in the initial data the job was created with.
+                            elif param_name in context.initial_data:
+                                params_to_inject[param_name] = context.initial_data[param_name]
+
+                    await handler(**params_to_inject)
+
+                    duration_ms = int((monotonic() - start_time) * 1000)
+
+                    # Process the single action requested by the handler.
+                    if action_factory.next_state:
+                        await self._handle_transition(
+                            job_state,
+                            action_factory.next_state,
+                            duration_ms,
+                        )
+                    elif action_factory.task_to_dispatch:
+                        await self._handle_dispatch(
+                            job_state,
+                            action_factory.task_to_dispatch,
+                            duration_ms,
+                        )
+                    elif action_factory.parallel_tasks_to_dispatch:
+                        await self._handle_parallel_dispatch(
+                            job_state,
+                            action_factory.parallel_tasks_to_dispatch,
+                            duration_ms,
+                        )
+                    elif action_factory.sub_blueprint_to_run:
+                        await self._handle_run_blueprint(
+                            job_state,
+                            action_factory.sub_blueprint_to_run,
+                            duration_ms,
+                        )
+
+                except Exception as e:
+                    # This catches errors within the handler's execution.
+                    duration_ms = int((monotonic() - start_time) * 1000)
+                    await self._handle_failure(job_state, e, duration_ms)
+        finally:
+            await self.storage.ack_job(message_id)
+            if message_id in self._processing_messages:
+                self._processing_messages.remove(message_id)
 
     async def _handle_transition(
         self,
-        job_state: Dict[str, Any],
+        job_state: dict[str, Any],
         next_state: str,
         duration_ms: int,
     ):
@@ -258,8 +263,8 @@ class JobExecutor:
 
     async def _handle_dispatch(
         self,
-        job_state: Dict[str, Any],
-        task_info: Dict[str, Any],
+        job_state: dict[str, Any],
+        task_info: dict[str, Any],
         duration_ms: int,
     ):
         job_id = job_state["id"]
@@ -302,8 +307,8 @@ class JobExecutor:
 
     async def _handle_run_blueprint(
         self,
-        parent_job_state: Dict[str, Any],
-        sub_blueprint_info: Dict[str, Any],
+        parent_job_state: dict[str, Any],
+        sub_blueprint_info: dict[str, Any],
         duration_ms: int,
     ):
         parent_job_id = parent_job_state["id"]
@@ -342,8 +347,8 @@ class JobExecutor:
 
     async def _handle_parallel_dispatch(
         self,
-        job_state: Dict[str, Any],
-        parallel_info: Dict[str, Any],
+        job_state: dict[str, Any],
+        parallel_info: dict[str, Any],
         duration_ms: int,
     ):
         job_id = job_state["id"]
@@ -390,7 +395,7 @@ class JobExecutor:
 
     async def _handle_failure(
         self,
-        job_state: Dict[str, Any],
+        job_state: dict[str, Any],
         error: Exception,
         duration_ms: int,
     ):
@@ -448,7 +453,7 @@ class JobExecutor:
                 {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]):
         """Checks if a completed job was a sub-job. If so, it resumes the parent
         job, passing the success/failure outcome of the child.
         """
@@ -501,14 +506,29 @@ class JobExecutor:
             logger.exception("Unhandled exception in job processing task")
 
     async def run(self):
+        import asyncio
+
         logger.info("JobExecutor started.")
         self._running = True
+        semaphore = asyncio.Semaphore(self.engine.config.EXECUTOR_MAX_CONCURRENT_JOBS)
+
         while self._running:
             try:
-                job_id = await self.storage.dequeue_job()
-                if job_id:
-                    task = create_task(self._process_job(job_id))
+                # Wait for an available slot before fetching a new job
+                await semaphore.acquire()
+
+                result = await self.storage.dequeue_job()
+                if result:
+                    job_id, message_id = result
+                    task = create_task(self._process_job(job_id, message_id))
                     task.add_done_callback(self._handle_task_completion)
+                    # 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
+                    semaphore.release()
+                    # Prevent busy loop if storage returns None immediately
+                    await sleep(0.1)
             except CancelledError:
                 break
             except Exception:

avtomatika/history/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
+from typing import Any
 
 
 class HistoryStorageBase(ABC):
@@ -13,29 +13,29 @@ class HistoryStorageBase(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def log_job_event(self, event_data: Dict[str, Any]):
+    async def log_job_event(self, event_data: dict[str, Any]):
         """Logs an event related to the job lifecycle."""
         raise NotImplementedError
 
     @abstractmethod
-    async def log_worker_event(self, event_data: Dict[str, Any]):
+    async def log_worker_event(self, event_data: dict[str, Any]):
         """Logs an event related to the worker lifecycle."""
         raise NotImplementedError
 
     @abstractmethod
-    async def get_job_history(self, job_id: str) -> List[Dict[str, Any]]:
+    async def get_job_history(self, job_id: str) -> list[dict[str, Any]]:
         """Gets the full history for the specified job."""
         raise NotImplementedError
 
     @abstractmethod
-    async def get_jobs(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
+    async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         """Gets a paginated list of recent jobs.
         Primarily returns the last event for each job.
         """
         raise NotImplementedError
 
     @abstractmethod
-    async def get_job_summary(self) -> Dict[str, int]:
+    async def get_job_summary(self) -> dict[str, int]:
         """Returns a summary of job statuses.
         Example: {'running': 10, 'completed': 50, 'failed': 5}
         """
@@ -46,6 +46,6 @@ class HistoryStorageBase(ABC):
         self,
         worker_id: str,
         since_days: int,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         """Gets the event history for a specific worker for the last N days."""
         raise NotImplementedError

avtomatika/history/noop.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any
 
 from .base import HistoryStorageBase
 
@@ -12,27 +12,27 @@ class NoOpHistoryStorage(HistoryStorageBase):
         # Do nothing
         pass
 
-    async def log_job_event(self, event_data: Dict[str, Any]):
+    async def log_job_event(self, event_data: dict[str, Any]):
         # Do nothing
         pass
 
-    async def log_worker_event(self, event_data: Dict[str, Any]):
+    async def log_worker_event(self, event_data: dict[str, Any]):
         # Do nothing
         pass
 
-    async def get_job_history(self, job_id: str) -> List[Dict[str, Any]]:
+    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]]:
+    async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:
         return []
 
-    async def get_job_summary(self) -> Dict[str, int]:
+    async def get_job_summary(self) -> dict[str, int]:
         return {}
 
     async def get_worker_history(
         self,
         worker_id: str,
         since_days: int,
-    ) -> List[Dict[str, Any]]:
+    ) -> list[dict[str, Any]]:
         return []