avtomatika 1.0b7__py3-none-any.whl → 1.0b9__py3-none-any.whl

avtomatika/services/worker_service.py

@@ -0,0 +1,266 @@
+from hashlib import sha256
+from logging import getLogger
+from secrets import token_urlsafe
+from time import monotonic
+from typing import Any, Optional
+
+from rxon.models import TokenResponse
+from rxon.validators import validate_identifier
+
+from ..app_keys import S3_SERVICE_KEY
+from ..config import Config
+from ..constants import (
+    ERROR_CODE_INTEGRITY_MISMATCH,
+    ERROR_CODE_INVALID_INPUT,
+    ERROR_CODE_PERMANENT,
+    ERROR_CODE_TRANSIENT,
+    JOB_STATUS_CANCELLED,
+    JOB_STATUS_FAILED,
+    JOB_STATUS_QUARANTINED,
+    JOB_STATUS_RUNNING,
+    JOB_STATUS_WAITING_FOR_PARALLEL,
+    TASK_STATUS_CANCELLED,
+    TASK_STATUS_FAILURE,
+    TASK_STATUS_SUCCESS,
+)
+from ..history.base import HistoryStorageBase
+from ..storage.base import StorageBackend
+
+logger = getLogger(__name__)
+
+
+class WorkerService:
+    def __init__(
+        self,
+        storage: StorageBackend,
+        history_storage: HistoryStorageBase,
+        config: Config,
+        engine: Any,
+    ):
+        self.storage = storage
+        self.history_storage = history_storage
+        self.config = config
+        self.engine = engine
+
+    async def register_worker(self, worker_data: dict[str, Any]) -> None:
+        """
+        Registers a new worker.
+        :param worker_data: Raw dictionary from request (to be validated/converted to Model later)
+        """
+        worker_id = worker_data.get("worker_id")
+        if not worker_id:
+            raise ValueError("Missing required field: worker_id")
+
+        validate_identifier(worker_id, "worker_id")
+
+        # S3 Consistency Check
+        s3_service = self.engine.app.get(S3_SERVICE_KEY)
+        if s3_service:
+            orchestrator_s3_hash = s3_service.get_config_hash()
+            worker_capabilities = worker_data.get("capabilities", {})
+            worker_s3_hash = worker_capabilities.get("s3_config_hash")
+
+            if orchestrator_s3_hash and worker_s3_hash and orchestrator_s3_hash != worker_s3_hash:
+                logger.warning(
+                    f"Worker '{worker_id}' has a different S3 configuration hash! "
+                    f"Orchestrator: {orchestrator_s3_hash}, Worker: {worker_s3_hash}. "
+                    "This may lead to 'split-brain' storage issues."
+                )
+
+        ttl = self.config.WORKER_HEALTH_CHECK_INTERVAL_SECONDS * 2
+        await self.storage.register_worker(worker_id, worker_data, ttl)
+
+        logger.info(f"Worker '{worker_id}' registered with info: {worker_data}")
+
+        await self.history_storage.log_worker_event(
+            {
+                "worker_id": worker_id,
+                "event_type": "registered",
+                "worker_info_snapshot": worker_data,
+            }
+        )
+
+    async def get_next_task(self, worker_id: str) -> Optional[dict[str, Any]]:
+        """
+        Retrieves the next task for a worker using long-polling configuration.
+        """
+        logger.debug(f"Worker {worker_id} is requesting a new task.")
+        return await self.storage.dequeue_task_for_worker(worker_id, self.config.WORKER_POLL_TIMEOUT_SECONDS)
+
+    async def process_task_result(self, result_payload: dict[str, Any], authenticated_worker_id: str) -> str:
+        """
+        Processes a task result submitted by a worker.
+        Returns a status string constant.
+        """
+        payload_worker_id = result_payload.get("worker_id")
+
+        if payload_worker_id and payload_worker_id != authenticated_worker_id:
+            raise PermissionError(
+                f"Forbidden: Authenticated worker '{authenticated_worker_id}' "
+                f"cannot submit results for another worker '{payload_worker_id}'."
+            )
+
+        job_id = result_payload.get("job_id")
+        task_id = result_payload.get("task_id")
+        result_data = result_payload.get("result", {})
+
+        if not job_id or not task_id:
+            raise ValueError("job_id and task_id are required")
+
+        job_state = await self.storage.get_job_state(job_id)
+        if not job_state:
+            raise LookupError("Job not found")
+
+        if job_state.get("status") == JOB_STATUS_WAITING_FOR_PARALLEL:
+            await self.storage.remove_job_from_watch(f"{job_id}:{task_id}")
+            job_state.setdefault("aggregation_results", {})[task_id] = result_data
+
+            branches = job_state.setdefault("active_branches", [])
+            if task_id in branches:
+                branches.remove(task_id)
+
+            if not branches:
+                logger.info(f"All parallel branches for job {job_id} have completed.")
+                job_state["status"] = JOB_STATUS_RUNNING
+                job_state["current_state"] = job_state["aggregation_target"]
+                await self.storage.save_job_state(job_id, job_state)
+                await self.storage.enqueue_job(job_id)
+            else:
+                logger.info(
+                    f"Branch {task_id} for job {job_id} completed. Waiting for {len(branches)} more.",
+                )
+                await self.storage.save_job_state(job_id, job_state)
+
+            return "parallel_branch_result_accepted"
+
+        await self.storage.remove_job_from_watch(job_id)
+
+        now = monotonic()
+        dispatched_at = job_state.get("task_dispatched_at", now)
+        duration_ms = int((now - dispatched_at) * 1000)
+
+        await self.history_storage.log_job_event(
+            {
+                "job_id": job_id,
+                "state": job_state.get("current_state"),
+                "event_type": "task_finished",
+                "duration_ms": duration_ms,
+                "worker_id": authenticated_worker_id,
+                "context_snapshot": {**job_state, "result": result_data},
+            },
+        )
+
+        result_status = result_data.get("status", TASK_STATUS_SUCCESS)  # Default to success? Constant?
+
+        if result_status == TASK_STATUS_FAILURE:
+            return await self._handle_task_failure(job_state, task_id, result_data)
+
+        if result_status == TASK_STATUS_CANCELLED:
+            logger.info(f"Task {task_id} for job {job_id} was cancelled by worker.")
+            job_state["status"] = JOB_STATUS_CANCELLED
+            await self.storage.save_job_state(job_id, job_state)
+
+            transitions = job_state.get("current_task_transitions", {})
+            if next_state := transitions.get("cancelled"):
+                job_state["current_state"] = next_state
+                job_state["status"] = JOB_STATUS_RUNNING
+                await self.storage.save_job_state(job_id, job_state)
+                await self.storage.enqueue_job(job_id)
+            return "result_accepted_cancelled"
+
+        transitions = job_state.get("current_task_transitions", {})
+        result_status = result_data.get("status", TASK_STATUS_SUCCESS)
+        next_state = transitions.get(result_status)
+
+        if next_state:
+            logger.info(f"Job {job_id} transitioning based on worker status '{result_status}' to state '{next_state}'")
+
+            worker_data_content = result_data.get("data")
+            if worker_data_content and isinstance(worker_data_content, dict):
+                if "state_history" not in job_state:
+                    job_state["state_history"] = {}
+                job_state["state_history"].update(worker_data_content)
+
+            data_metadata = result_payload.get("data_metadata")
+            if data_metadata:
+                if "data_metadata" not in job_state:
+                    job_state["data_metadata"] = {}
+                job_state["data_metadata"].update(data_metadata)
+                logger.debug(f"Stored data metadata for job {job_id}: {list(data_metadata.keys())}")
+
+            job_state["current_state"] = next_state
+            job_state["status"] = JOB_STATUS_RUNNING
+            await self.storage.save_job_state(job_id, job_state)
+            await self.storage.enqueue_job(job_id)
+            return "result_accepted_success"
+        else:
+            logger.error(f"Job {job_id} failed. Worker returned unhandled status '{result_status}'.")
+            job_state["status"] = JOB_STATUS_FAILED
+            job_state["error_message"] = f"Worker returned unhandled status: {result_status}"
+            await self.storage.save_job_state(job_id, job_state)
+            return "result_accepted_failure"
+
+    async def _handle_task_failure(self, job_state: dict, task_id: str, result_data: dict) -> str:
+        error_details = result_data.get("error", {})
+        error_type = ERROR_CODE_TRANSIENT
+        error_message = "No error details provided."
+
+        if isinstance(error_details, dict):
+            error_type = error_details.get("code", ERROR_CODE_TRANSIENT)
+            error_message = error_details.get("message", "No error message provided.")
+        elif isinstance(error_details, str):
+            error_message = error_details
+
+        job_id = job_state["id"]
+        logger.warning(f"Task {task_id} for job {job_id} failed with error type '{error_type}'.")
+
+        if error_type == ERROR_CODE_PERMANENT:
+            job_state["status"] = JOB_STATUS_QUARANTINED
+            job_state["error_message"] = f"Task failed with permanent error: {error_message}"
+            await self.storage.save_job_state(job_id, job_state)
+            await self.storage.quarantine_job(job_id)
+        elif error_type == ERROR_CODE_INVALID_INPUT:
+            job_state["status"] = JOB_STATUS_FAILED
+            job_state["error_message"] = f"Task failed due to invalid input: {error_message}"
+            await self.storage.save_job_state(job_id, job_state)
+        elif error_type == ERROR_CODE_INTEGRITY_MISMATCH:
+            job_state["status"] = JOB_STATUS_FAILED
+            job_state["error_message"] = f"Task failed due to data integrity mismatch: {error_message}"
+            await self.storage.save_job_state(job_id, job_state)
+            logger.critical(f"Data integrity mismatch detected for job {job_id}: {error_message}")
+        else:
+            await self.engine.handle_task_failure(job_state, task_id, error_message)
+
+        return "result_accepted_failure"
+
+    async def issue_access_token(self, worker_id: str) -> TokenResponse:
+        """Generates and stores a temporary access token."""
+        raw_token = token_urlsafe(32)
+        token_hash = sha256(raw_token.encode()).hexdigest()
+        ttl = 3600
+
+        await self.storage.save_worker_access_token(worker_id, token_hash, ttl)
+        logger.info(f"Issued temporary access token for worker {worker_id}")
+
+        return TokenResponse(access_token=raw_token, expires_in=ttl, worker_id=worker_id)
+
+    async def update_worker_heartbeat(
+        self, worker_id: str, update_data: Optional[dict[str, Any]]
+    ) -> Optional[dict[str, Any]]:
+        """Updates worker TTL and status."""
+        ttl = self.config.WORKER_HEALTH_CHECK_INTERVAL_SECONDS * 2
+
+        if update_data:
+            updated_worker = await self.storage.update_worker_status(worker_id, update_data, ttl)
+            if updated_worker:
+                await self.history_storage.log_worker_event(
+                    {
+                        "worker_id": worker_id,
+                        "event_type": "status_update",
+                        "worker_info_snapshot": updated_worker,
+                    },
+                )
+            return updated_worker
+        else:
+            refreshed = await self.storage.refresh_worker_ttl(worker_id, ttl)
+            return {"status": "ttl_refreshed"} if refreshed else None

avtomatika/storage/base.py

@@ -142,6 +142,37 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    async def get_active_worker_ids(self) -> list[str]:
+        """Returns a list of IDs for all currently active workers.
+
+        :return: A list of worker ID strings.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def cleanup_expired_workers(self) -> None:
+        """Maintenance task to clean up internal indexes from expired worker entries."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get_workers(self, worker_ids: list[str]) -> list[dict[str, Any]]:
+        """Bulk retrieves worker info for a list of IDs.
+
+        :param worker_ids: List of worker identifiers.
+        :return: List of worker info dictionaries.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def find_workers_for_task(self, task_type: str) -> list[str]:
+        """Finds idle workers that support the given task.
+
+        :param task_type: The type of task to find workers for.
+        :return: A list of worker IDs that are idle and support the task.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     async def add_job_to_watch(self, job_id: str, timeout_at: float) -> None:
         """Add a job to the list for timeout tracking.
@@ -152,9 +183,10 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_timed_out_jobs(self) -> list[str]:
+    async def get_timed_out_jobs(self, limit: int = 100) -> list[str]:
         """Get a list of job IDs that are overdue and remove them from the tracking list.
 
+        :param limit: Maximum number of jobs to retrieve.
         :return: A list of overdue job IDs.
         """
         raise NotImplementedError
@@ -165,9 +197,10 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def dequeue_job(self) -> tuple[str, str] | None:
+    async def dequeue_job(self, block: int | None = None) -> tuple[str, str] | None:
         """Retrieve a job ID and its message ID from the execution queue.
 
+        :param block: Milliseconds to block if no message is available. None for non-blocking.
         :return: A tuple of (job_id, message_id) or None if the timeout has expired.
         """
         raise NotImplementedError
@@ -250,7 +283,7 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def set_worker_token(self, worker_id: str, token: str):
+    async def set_worker_token(self, worker_id: str, token: str) -> None:
         """Saves an individual token for a specific worker."""
         raise NotImplementedError
 
@@ -259,13 +292,23 @@ class StorageBackend(ABC):
         """Retrieves an individual token for a specific worker."""
         raise NotImplementedError
 
+    @abstractmethod
+    async def save_worker_access_token(self, worker_id: str, token: str, ttl: int) -> None:
+        """Saves a temporary access token for a worker (STS)."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def verify_worker_access_token(self, token: str) -> str | None:
+        """Verifies a temporary access token and returns the associated worker_id if valid."""
+        raise NotImplementedError
+
     @abstractmethod
     async def get_worker_info(self, worker_id: str) -> dict[str, Any] | None:
         """Get complete information about a worker by its ID."""
         raise NotImplementedError
 
     @abstractmethod
-    async def flush_all(self):
+    async def flush_all(self) -> None:
         """Completely clears the storage. Used mainly for tests."""
         raise NotImplementedError
 
@@ -312,3 +355,11 @@ class StorageBackend(ABC):
         :return: True if the lock was successfully released, False otherwise.
         """
         raise NotImplementedError
+
+    @abstractmethod
+    async def ping(self) -> bool:
+        """Checks connection to the storage backend.
+
+        :return: True if storage is accessible, False otherwise.
+        """
+        raise NotImplementedError

avtomatika/storage/memory.py

@@ -154,6 +154,33 @@ class MemoryStorage(StorageBackend):
             )
             return active_workers
 
+    async def get_active_worker_ids(self) -> list[str]:
+        async with self._lock:
+            now = monotonic()
+            return [worker_id for worker_id, ttl in self._worker_ttls.items() if ttl > now]
+
+    async def cleanup_expired_workers(self) -> None:
+        async with self._lock:
+            await self._clean_expired()
+
+    async def get_workers(self, worker_ids: list[str]) -> list[dict[str, Any]]:
+        async with self._lock:
+            return [self._workers[wid] for wid in worker_ids if wid in self._workers]
+
+    async def find_workers_for_task(self, task_type: str) -> list[str]:
+        """Finds idle workers supporting the task (O(N) for memory storage)."""
+        async with self._lock:
+            now = monotonic()
+            candidates = []
+            for worker_id, info in self._workers.items():
+                if self._worker_ttls.get(worker_id, 0) <= now:
+                    continue
+                if info.get("status", "idle") != "idle":
+                    continue
+                if task_type in info.get("supported_tasks", []):
+                    candidates.append(worker_id)
+            return candidates
+
     async def add_job_to_watch(self, job_id: str, timeout_at: float) -> None:
         async with self._lock:
             self._watched_jobs[job_id] = timeout_at
@@ -162,10 +189,11 @@ class MemoryStorage(StorageBackend):
         async with self._lock:
             self._watched_jobs.pop(job_id, None)
 
-    async def get_timed_out_jobs(self) -> list[str]:
+    async def get_timed_out_jobs(self, limit: int = 100) -> list[str]:
         async with self._lock:
             now = monotonic()
             timed_out_ids = [job_id for job_id, timeout_at in self._watched_jobs.items() if timeout_at <= now]
+            timed_out_ids = timed_out_ids[:limit]
             for job_id in timed_out_ids:
                 self._watched_jobs.pop(job_id, None)
             return timed_out_ids
@@ -173,13 +201,21 @@ class MemoryStorage(StorageBackend):
     async def enqueue_job(self, job_id: str) -> None:
         await self._job_queue.put(job_id)
 
-    async def dequeue_job(self) -> tuple[str, str] | None:
-        """Waits indefinitely for a job ID from the queue and returns it.
-        Returns a tuple of (job_id, message_id). In MemoryStorage, message_id is dummy.
+    async def dequeue_job(self, block: int | None = None) -> tuple[str, str] | None:
+        """Waits for a job ID from the queue.
+        If block is None, waits indefinitely.
+        If block is int, waits for that many milliseconds.
         """
-        job_id = await self._job_queue.get()
-        self._job_queue.task_done()
-        return job_id, "memory-msg-id"
+        try:
+            if block is None:
+                job_id = await self._job_queue.get()
+            else:
+                job_id = await wait_for(self._job_queue.get(), timeout=block / 1000.0)
+
+            self._job_queue.task_done()
+            return job_id, "memory-msg-id"
+        except AsyncTimeoutError:
+            return None
 
     async def ack_job(self, message_id: str) -> None:
         """No-op for MemoryStorage as it doesn't support persistent streams."""
@@ -296,6 +332,16 @@ class MemoryStorage(StorageBackend):
         async with self._lock:
             return self._worker_tokens.get(worker_id)
 
+    async def save_worker_access_token(self, worker_id: str, token: str, ttl: int) -> None:
+        async with self._lock:
+            self._generic_keys[f"sts:{token}"] = worker_id
+            self._generic_key_ttls[f"sts:{token}"] = monotonic() + ttl
+
+    async def verify_worker_access_token(self, token: str) -> str | None:
+        async with self._lock:
+            await self._clean_expired()
+            return self._generic_keys.get(f"sts:{token}")
+
     async def set_task_cancellation_flag(self, task_id: str) -> None:
         key = f"task_cancel:{task_id}"
         await self.increment_key_with_ttl(key, 3600)
@@ -334,3 +380,6 @@ class MemoryStorage(StorageBackend):
                     del self._locks[key]
                     return True
             return False
+
+    async def ping(self) -> bool:
+        return True