avtomatika 1.0b2__tar.gz → 1.0b3__tar.gz

{avtomatika-1.0b2/src/avtomatika.egg-info → avtomatika-1.0b3}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: avtomatika
-Version: 1.0b2
-Summary: A state-machine based orchestrator for long-running jobs.
+Version: 1.0b3
+Summary: A state-machine based orchestrator for long-running AI and other jobs.
 Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika
 Project-URL: Bug Tracker, https://github.com/avtomatika-ai/avtomatika/issues
 Classifier: Development Status :: 4 - Beta
@@ -18,25 +18,25 @@ Requires-Dist: graphviz~=0.21
 Requires-Dist: zstandard~=0.24
 Requires-Dist: aioprometheus~=23.12
 Provides-Extra: redis
-Requires-Dist: redis~=6.4; extra == "redis"
+Requires-Dist: redis~=7.1; extra == "redis"
 Requires-Dist: orjson~=3.11; extra == "redis"
 Provides-Extra: history
-Requires-Dist: aiosqlite~=0.21; extra == "history"
+Requires-Dist: aiosqlite~=0.22; extra == "history"
 Requires-Dist: asyncpg~=0.30; extra == "history"
 Requires-Dist: orjson~=3.11; extra == "history"
 Provides-Extra: telemetry
-Requires-Dist: opentelemetry-api~=1.38; extra == "telemetry"
-Requires-Dist: opentelemetry-sdk~=1.38; extra == "telemetry"
-Requires-Dist: opentelemetry-exporter-otlp~=1.36; extra == "telemetry"
+Requires-Dist: opentelemetry-api~=1.39; extra == "telemetry"
+Requires-Dist: opentelemetry-sdk~=1.39; extra == "telemetry"
+Requires-Dist: opentelemetry-exporter-otlp~=1.39; extra == "telemetry"
 Requires-Dist: opentelemetry-instrumentation-aiohttp-client~=0.59b0; extra == "telemetry"
 Provides-Extra: test
-Requires-Dist: pytest~=8.4; extra == "test"
+Requires-Dist: pytest~=9.0; extra == "test"
 Requires-Dist: pytest-asyncio~=1.1; extra == "test"
-Requires-Dist: fakeredis~=2.31; extra == "test"
+Requires-Dist: fakeredis~=2.33; extra == "test"
 Requires-Dist: pytest-aiohttp~=1.1; extra == "test"
 Requires-Dist: pytest-mock~=3.14; extra == "test"
 Requires-Dist: aioresponses~=0.7; extra == "test"
-Requires-Dist: backports.zstd; extra == "test"
+Requires-Dist: backports.zstd~=1.2; extra == "test"
 Requires-Dist: opentelemetry-instrumentation-aiohttp-client; extra == "test"
 Provides-Extra: all
 Requires-Dist: avtomatika[redis]; extra == "all"
@@ -285,7 +285,7 @@ Run multiple tasks simultaneously and gather their results.
 @my_blueprint.handler_for("process_files")
 async def fan_out_handler(initial_data, actions):
     tasks_to_dispatch = [
-        {"task_type": "file_analysis", "params": {"file": file}}
+        {"task_type": "file_analysis", "params": {"file": file}})
         for file in initial_data.get("files", [])
     ]
     # Use dispatch_parallel to send all tasks at once.
@@ -332,6 +332,8 @@ async def cache_handler(data_stores):
 
 The orchestrator's behavior can be configured through environment variables. Additionally, any configuration parameter loaded from environment variables can be programmatically overridden in your application code after the `Config` object has been initialized. This provides flexibility for different deployment and testing scenarios.
 
+**Important:** The system employs **strict validation** for configuration files (`clients.toml`, `workers.toml`) at startup. If a configuration file is invalid (e.g., malformed TOML, missing required fields), the application will **fail fast** and exit with an error, rather than starting in a partially broken state. This ensures the security and integrity of the deployment.
+
 ### Fault Tolerance
 
 The orchestrator has built-in mechanisms for handling failures based on the `error.code` field in a worker's response.
@@ -340,6 +342,13 @@ The orchestrator has built-in mechanisms for handling failures based on the `err
 *   **PERMANENT_ERROR**: A permanent error (e.g., a corrupted file). The task will be immediately sent to quarantine for manual investigation.
 *   **INVALID_INPUT_ERROR**: An error in the input data. The entire pipeline (Job) will be immediately moved to the failed state.
 
+### High Availability & Distributed Locking
+
+The architecture supports horizontal scaling. Multiple Orchestrator instances can run behind a load balancer.
+
+*   **Stateless API:** The API is stateless; all state is persisted in Redis.
+*   **Distributed Locking:** Background processes (`Watcher`, `ReputationCalculator`) use distributed locks (via Redis `SET NX`) to coordinate and prevent race conditions when multiple instances are active.
+
 ### Storage Backend
 
 By default, the engine uses in-memory storage. For production, you must configure persistent storage via environment variables.
@@ -408,3 +417,12 @@ To run the `avtomatika` test suite:
 ```bash
 pytest avtomatika/tests/
 ```
+
+## Detailed Documentation
+
+For a deeper dive into the system, please refer to the following documents in the `docs/` directory:
+
+- [**Architecture Guide**](docs/architecture.md): A detailed overview of the system components and their interactions.
+- [**API Reference**](docs/api_reference.md): Full specification of the HTTP API.
+- [**Deployment Guide**](docs/deployment.md): Instructions for deploying with Gunicorn/Uvicorn and NGINX.
+- [**Cookbook**](docs/cookbook/README.md): Examples and best practices for creating blueprints.

{avtomatika-1.0b2 → avtomatika-1.0b3}/README.md

@@ -239,7 +239,7 @@ Run multiple tasks simultaneously and gather their results.
 @my_blueprint.handler_for("process_files")
 async def fan_out_handler(initial_data, actions):
     tasks_to_dispatch = [
-        {"task_type": "file_analysis", "params": {"file": file}}
+        {"task_type": "file_analysis", "params": {"file": file}})
         for file in initial_data.get("files", [])
     ]
     # Use dispatch_parallel to send all tasks at once.
@@ -286,6 +286,8 @@ async def cache_handler(data_stores):
 
 The orchestrator's behavior can be configured through environment variables. Additionally, any configuration parameter loaded from environment variables can be programmatically overridden in your application code after the `Config` object has been initialized. This provides flexibility for different deployment and testing scenarios.
 
+**Important:** The system employs **strict validation** for configuration files (`clients.toml`, `workers.toml`) at startup. If a configuration file is invalid (e.g., malformed TOML, missing required fields), the application will **fail fast** and exit with an error, rather than starting in a partially broken state. This ensures the security and integrity of the deployment.
+
 ### Fault Tolerance
 
 The orchestrator has built-in mechanisms for handling failures based on the `error.code` field in a worker's response.
@@ -294,6 +296,13 @@ The orchestrator has built-in mechanisms for handling failures based on the `err
 *   **PERMANENT_ERROR**: A permanent error (e.g., a corrupted file). The task will be immediately sent to quarantine for manual investigation.
 *   **INVALID_INPUT_ERROR**: An error in the input data. The entire pipeline (Job) will be immediately moved to the failed state.
 
+### High Availability & Distributed Locking
+
+The architecture supports horizontal scaling. Multiple Orchestrator instances can run behind a load balancer.
+
+*   **Stateless API:** The API is stateless; all state is persisted in Redis.
+*   **Distributed Locking:** Background processes (`Watcher`, `ReputationCalculator`) use distributed locks (via Redis `SET NX`) to coordinate and prevent race conditions when multiple instances are active.
+
 ### Storage Backend
 
 By default, the engine uses in-memory storage. For production, you must configure persistent storage via environment variables.
@@ -362,3 +371,12 @@ To run the `avtomatika` test suite:
 ```bash
 pytest avtomatika/tests/
 ```
+
+## Detailed Documentation
+
+For a deeper dive into the system, please refer to the following documents in the `docs/` directory:
+
+- [**Architecture Guide**](docs/architecture.md): A detailed overview of the system components and their interactions.
+- [**API Reference**](docs/api_reference.md): Full specification of the HTTP API.
+- [**Deployment Guide**](docs/deployment.md): Instructions for deploying with Gunicorn/Uvicorn and NGINX.
+- [**Cookbook**](docs/cookbook/README.md): Examples and best practices for creating blueprints.

{avtomatika-1.0b2 → avtomatika-1.0b3}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "avtomatika"
-version = "1.0b2"
-description = "A state-machine based orchestrator for long-running jobs."
+version = "1.0b3"
+description = "A state-machine based orchestrator for long-running AI and other jobs."
 readme = "README.md"
 requires-python = ">=3.11"
 classifiers = [
@@ -24,22 +24,22 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-redis = ["redis~=6.4", "orjson~=3.11"]
-history = ["aiosqlite~=0.21", "asyncpg~=0.30", "orjson~=3.11"]
+redis = ["redis~=7.1", "orjson~=3.11"]
+history = ["aiosqlite~=0.22", "asyncpg~=0.30", "orjson~=3.11"]
 telemetry = [
-    "opentelemetry-api~=1.38",
-    "opentelemetry-sdk~=1.38",
-    "opentelemetry-exporter-otlp~=1.36",
+    "opentelemetry-api~=1.39",
+    "opentelemetry-sdk~=1.39",
+    "opentelemetry-exporter-otlp~=1.39",
     "opentelemetry-instrumentation-aiohttp-client~=0.59b0",
 ]
 test = [
-    "pytest~=8.4",
+    "pytest~=9.0",
     "pytest-asyncio~=1.1",
-    "fakeredis~=2.31",
+    "fakeredis~=2.33",
     "pytest-aiohttp~=1.1",
     "pytest-mock~=3.14",
     "aioresponses~=0.7",
-    "backports.zstd",
+    "backports.zstd~=1.2",
     "opentelemetry-instrumentation-aiohttp-client",
 ]
 all = [

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/__init__.py

@@ -4,6 +4,7 @@
 This module exposes the primary classes for building and running state-driven automations.
 """
 
+import contextlib
 from importlib.metadata import version
 
 __version__ = version("avtomatika")
@@ -22,9 +23,7 @@ __all__ = [
     "StorageBackend",
 ]
 
-try:
+with contextlib.suppress(ImportError):
     from .storage.redis import RedisStorage  # noqa: F401
 
     __all__.append("RedisStorage")
-except ImportError:
-    pass

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/blueprint.py

@@ -168,8 +168,7 @@ class StateMachineBlueprint:
         for handler in self.conditional_handlers:
             if handler.state == state and handler.evaluate(context):
                 return handler.func
-        default_handler = self.handlers.get(state)
-        if default_handler:
+        if default_handler := self.handlers.get(state):
             return default_handler
         raise ValueError(
             f"No suitable handler found for state '{state}' in blueprint '{self.name}' for the given context.",
@@ -230,12 +229,11 @@ class StateMachineBlueprint:
                     f"Could not parse handler '{handler_func.__name__}' for state '{handler_state}'. "
                     f"Graph may be incomplete. Error: {e}"
                 )
-                pass
         for state in states:
             dot.node(state, state)
 
-        if output_filename:
-            dot.render(output_filename, format=output_format, cleanup=True)
-            print(f"Graph rendered to {output_filename}.{output_format}")
-        else:
+        if not output_filename:
             return dot.source
+        dot.render(output_filename, format=output_format, cleanup=True)
+        print(f"Graph rendered to {output_filename}.{output_format}")
+        return None

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/client_config_loader.py

@@ -26,25 +26,37 @@ async def load_client_configs_to_redis(
             config_path,
         )
         return
+    except Exception as e:
+        logger.error(f"Failed to parse client config file '{config_path}': {e}")
+        raise ValueError(f"Invalid client configuration file: {e}") from e
 
     loaded_count = 0
     for client_name, config in clients_data.items():
+        if not isinstance(config, dict):
+            logger.error(f"Client '{client_name}' configuration must be a table (dict).")
+            raise ValueError(f"Invalid configuration for client '{client_name}'")
+
         token = config.get("token")
         if not token:
-            logger.warning(
-                "Skipping client '%s' due to missing 'token' field.",
-                client_name,
-            )
-            continue
+            logger.error(f"Client '{client_name}' is missing required 'token' field.")
+            raise ValueError(f"Missing token for client '{client_name}'")
+
+        if not isinstance(token, str):
+            logger.error(f"Token for client '{client_name}' must be a string.")
+            raise ValueError(f"Invalid token type for client '{client_name}'")
 
         # Separate static config from dynamic quota values
         static_config = {k: v for k, v in config.items() if k != "monthly_attempts"}
         quota = config.get("monthly_attempts")
 
+        if quota is not None and not isinstance(quota, int):
+            logger.error(f"Quota 'monthly_attempts' for client '{client_name}' must be an integer.")
+            raise ValueError(f"Invalid quota type for client '{client_name}'")
+
         try:
             # Assume these storage methods will be implemented
             await storage.save_client_config(token, static_config)
-            if quota is not None and isinstance(quota, int):
+            if quota is not None:
                 await storage.initialize_client_quota(token, quota)
 
             loaded_count += 1

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/dispatcher.py

@@ -28,15 +28,13 @@ class Dispatcher:
         self.config = config
         self._round_robin_indices: Dict[str, int] = defaultdict(int)
 
+    @staticmethod
     def _is_worker_compliant(
-        self,
         worker: Dict[str, Any],
         requirements: Dict[str, Any],
     ) -> bool:
         """Checks if a worker meets the specified resource requirements."""
-        # GPU check
-        required_gpu = requirements.get("gpu_info")
-        if required_gpu:
+        if required_gpu := requirements.get("gpu_info"):
             gpu_info = worker.get("resources", {}).get("gpu_info")
             if not gpu_info:
                 return False
@@ -51,17 +49,15 @@ class Dispatcher:
             ):
                 return False
 
-        # Installed models check
-        required_models = requirements.get("installed_models")
-        if required_models:
+        if required_models := requirements.get("installed_models"):
             installed_models = {m["name"] for m in worker.get("installed_models", [])}
             if not set(required_models).issubset(installed_models):
                 return False
 
         return True
 
+    @staticmethod
     def _select_default(
-        self,
         workers: List[Dict[str, Any]],
         task_type: str,
     ) -> Dict[str, Any]:
@@ -74,7 +70,7 @@ class Dispatcher:
         """
         warm_workers = [w for w in workers if task_type in w.get("hot_cache", [])]
 
-        target_pool = warm_workers if warm_workers else workers
+        target_pool = warm_workers or workers
 
         # The `cost` field is deprecated but maintained for backward compatibility.
         min_cost = min(w.get("cost", float("inf")) for w in target_pool)
@@ -95,8 +91,8 @@ class Dispatcher:
         self._round_robin_indices[task_type] = idx + 1
         return selected_worker
 
+    @staticmethod
     def _select_least_connections(
-        self,
         workers: List[Dict[str, Any]],
         task_type: str,
     ) -> Dict[str, Any]:
@@ -105,15 +101,16 @@ class Dispatcher:
         """
         return min(workers, key=lambda w: w.get("load", 0.0))
 
+    @staticmethod
     def _select_cheapest(
-        self,
         workers: List[Dict[str, Any]],
         task_type: str,
     ) -> Dict[str, Any]:
         """Selects the cheapest worker based on 'cost_per_second'."""
         return min(workers, key=lambda w: w.get("cost_per_second", float("inf")))
 
-    def _get_best_value_score(self, worker: Dict[str, Any]) -> float:
+    @staticmethod
+    def _get_best_value_score(worker: Dict[str, Any]) -> float:
         """Calculates a "score" for a worker using the formula cost / reputation.
         The lower the score, the better.
         """
@@ -121,9 +118,7 @@ class Dispatcher:
         # Default reputation is 1.0 if absent
         reputation = worker.get("reputation", 1.0)
         # Avoid division by zero
-        if reputation == 0:
-            return float("inf")
-        return cost / reputation
+        return float("inf") if reputation == 0 else cost / reputation
 
     def _select_best_value(
         self,
@@ -153,10 +148,9 @@ class Dispatcher:
         idle_workers = [w for w in all_workers if w.get("status", "idle") == "idle"]
         logger.debug(f"Idle workers: {[w['worker_id'] for w in idle_workers]}")
         if not idle_workers:
-            # If there are no idle workers, check if there are any busy workers in multi-orchestrator mode.
-            # This doesn't change the logic (an error will still occur), but it makes the logs more informative.
-            busy_mo_workers = [w for w in all_workers if w.get("status") == "busy" and "multi_orchestrator_info" in w]
-            if busy_mo_workers:
+            if busy_mo_workers := [
+                w for w in all_workers if w.get("status") == "busy" and "multi_orchestrator_info" in w
+            ]:
                 logger.warning(
                     f"No idle workers. Found {len(busy_mo_workers)} busy workers "
                     f"in multi-orchestrator mode. They are likely performing tasks for other Orchestrators.",

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/engine.py

@@ -485,8 +485,7 @@ class OrchestratorEngine:
             await self.storage.save_job_state(job_id, job_state)
             # Optionally, trigger a specific 'cancelled' transition if defined in the blueprint
             transitions = job_state.get("current_task_transitions", {})
-            next_state = transitions.get("cancelled")
-            if next_state:
+            if next_state := transitions.get("cancelled"):
                 job_state["current_state"] = next_state
                 job_state["status"] = "running"  # It's running the cancellation handler now
                 await self.storage.save_job_state(job_id, job_state)
@@ -494,9 +493,7 @@ class OrchestratorEngine:
             return web.json_response({"status": "result_accepted_cancelled"}, status=200)
 
         transitions = job_state.get("current_task_transitions", {})
-        next_state = transitions.get(result_status)
-
-        if next_state:
+        if next_state := transitions.get(result_status):
             logging.info(f"Job {job_id} transitioning based on worker status '{result_status}' to state '{next_state}'")
 
             worker_data = result.get("data")
@@ -602,7 +599,8 @@ class OrchestratorEngine:
         await load_client_configs_to_redis(self.storage)
         return web.json_response({"status": "db_flushed"}, status=200)
 
-    async def _docs_handler(self, request: web.Request) -> web.Response:
+    @staticmethod
+    async def _docs_handler(request: web.Request) -> web.Response:
         from importlib import resources
 
         try:
@@ -647,16 +645,7 @@ class OrchestratorEngine:
             all_protected_apps.append(protected_app)
 
         for app in all_protected_apps:
-            app.router.add_get("/jobs/{job_id}", self._get_job_status_handler)
-            app.router.add_post("/jobs/{job_id}/cancel", self._cancel_job_handler)
-            if not isinstance(self.history_storage, NoOpHistoryStorage):
-                app.router.add_get("/jobs/{job_id}/history", self._get_job_history_handler)
-            app.router.add_get("/blueprints/{blueprint_name}/graph", self._get_blueprint_graph_handler)
-            app.router.add_get("/workers", self._get_workers_handler)
-            app.router.add_get("/jobs", self._get_jobs_handler)
-            app.router.add_get("/dashboard", self._get_dashboard_handler)
-            app.router.add_post("/admin/reload-workers", self._reload_worker_configs_handler)
-
+            self._register_common_routes(app)
         if has_unversioned_routes:
             self.app.add_subapp("/api/", protected_app)
         for version, app in versioned_apps.items():
@@ -676,6 +665,17 @@ class OrchestratorEngine:
         worker_app.router.add_get("/ws/{worker_id}", self._websocket_handler)
         self.app.add_subapp("/_worker/", worker_app)
 
+    def _register_common_routes(self, app):
+        app.router.add_get("/jobs/{job_id}", self._get_job_status_handler)
+        app.router.add_post("/jobs/{job_id}/cancel", self._cancel_job_handler)
+        if not isinstance(self.history_storage, NoOpHistoryStorage):
+            app.router.add_get("/jobs/{job_id}/history", self._get_job_history_handler)
+        app.router.add_get("/blueprints/{blueprint_name}/graph", self._get_blueprint_graph_handler)
+        app.router.add_get("/workers", self._get_workers_handler)
+        app.router.add_get("/jobs", self._get_jobs_handler)
+        app.router.add_get("/dashboard", self._get_dashboard_handler)
+        app.router.add_post("/admin/reload-workers", self._reload_worker_configs_handler)
+
     async def _websocket_handler(self, request: web.Request) -> web.WebSocketResponse:
         worker_id = request.match_info.get("worker_id")
         if not worker_id:

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/executor.py

@@ -35,11 +35,13 @@ except ImportError:
         def inject(self, *args, **kwargs):
             pass
 
-        def extract(self, *args, **kwargs):
+        @staticmethod
+        def extract(*args, **kwargs):
             return None
 
     class NoOpTraceContextTextMapPropagator:
-        def extract(self, *args, **kwargs):
+        @staticmethod
+        def extract(*args, **kwargs):
             return None
 
     trace = NoOpTracer()
@@ -485,7 +487,8 @@ class JobExecutor:
         await self.storage.save_job_state(parent_job_id, parent_job_state)
         await self.storage.enqueue_job(parent_job_id)
 
-    def _handle_task_completion(self, task: Task):
+    @staticmethod
+    def _handle_task_completion(task: Task):
         """Callback to handle completion of a job processing task."""
         try:
             # This will re-raise any exception caught in the task

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/ratelimit.py

@@ -1,3 +1,4 @@
+from contextlib import suppress
 from typing import Awaitable, Callable
 
 from aiohttp import web
@@ -23,23 +24,15 @@ def rate_limit_middleware_factory(
         """Rate-limiting middleware that uses the provided storage backend."""
         # Determine the key for rate limiting (e.g., by worker_id or IP)
         # For worker endpoints, we key by worker_id. For others, by IP.
-        key_identifier = request.match_info.get("worker_id", request.remote)
-        if not key_identifier:
-            # Fallback for cases where remote IP might not be available
-            key_identifier = "unknown"
+        key_identifier = request.match_info.get("worker_id", request.remote) or "unknown"
 
         # Key by identifier and path to have per-endpoint limits
         rate_limit_key = f"ratelimit:{key_identifier}:{request.path}"
 
-        try:
+        with suppress(Exception):
             count = await storage.increment_key_with_ttl(rate_limit_key, period)
             if count > limit:
                 return web.json_response({"error": "Too Many Requests"}, status=429)
-        except Exception:
-            # If the rate limiter fails for any reason (e.g., Redis down),
-            # it's safer to let the request through than to block everything.
-            pass
-
         return await handler(request)
 
     return rate_limit_middleware

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/reputation.py

@@ -1,6 +1,7 @@
 from asyncio import CancelledError, sleep
 from logging import getLogger
 from typing import TYPE_CHECKING
+from uuid import uuid4
 
 if TYPE_CHECKING:
     from .engine import OrchestratorEngine
@@ -20,14 +21,22 @@ class ReputationCalculator:
         self.history_storage = engine.history_storage
         self.interval_seconds = interval_seconds
         self._running = False
+        self._instance_id = str(uuid4())
 
     async def run(self):
         """The main loop that periodically triggers reputation recalculation."""
-        logger.info("ReputationCalculator started.")
+        logger.info(f"ReputationCalculator started (Instance ID: {self._instance_id}).")
         self._running = True
         while self._running:
             try:
-                await self.calculate_all_reputations()
+                # Attempt to acquire lock
+                if await self.storage.acquire_lock("global_reputation_lock", self._instance_id, 300):
+                    try:
+                        await self.calculate_all_reputations()
+                    finally:
+                        await self.storage.release_lock("global_reputation_lock", self._instance_id)
+                else:
+                    logger.debug("ReputationCalculator lock held by another instance. Skipping.")
             except CancelledError:
                 break
             except Exception:

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/storage/__init__.py

@@ -1,11 +1,11 @@
+import contextlib
+
 from .base import StorageBackend
 from .memory import MemoryStorage
 
 __all__ = ["StorageBackend", "MemoryStorage"]
 
-try:
+with contextlib.suppress(ImportError):
     from .redis import RedisStorage  # noqa: F401
 
     __all__.append("RedisStorage")
-except ImportError:
-    pass

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/storage/base.py

@@ -264,3 +264,26 @@ class StorageBackend(ABC):
         Used for metrics.
         """
         raise NotImplementedError
+
+    @abstractmethod
+    async def acquire_lock(self, key: str, holder_id: str, ttl: int) -> bool:
+        """
+        Attempts to acquire a distributed lock.
+
+        :param key: The unique key of the lock (e.g., 'watcher_lock').
+        :param holder_id: A unique identifier for the caller (e.g., UUID).
+        :param ttl: Time-to-live for the lock in seconds.
+        :return: True if the lock was acquired, False otherwise.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def release_lock(self, key: str, holder_id: str) -> bool:
+        """
+        Releases a distributed lock if it is held by the specified holder_id.
+
+        :param key: The unique key of the lock.
+        :param holder_id: The identifier of the caller who presumably holds the lock.
+        :return: True if the lock was successfully released, False otherwise.
+        """
+        raise NotImplementedError

{avtomatika-1.0b2 → avtomatika-1.0b3}/src/avtomatika/storage/memory.py

@@ -25,6 +25,7 @@ class MemoryStorage(StorageBackend):
         self._worker_tokens: Dict[str, str] = {}
         self._generic_keys: Dict[str, Any] = {}
         self._generic_key_ttls: Dict[str, float] = {}
+        self._locks: Dict[str, tuple[str, float]] = {}  # key -> (holder_id, expiry_time)
 
         self._lock = Lock()
 
@@ -128,9 +129,11 @@ class MemoryStorage(StorageBackend):
         async with self._lock:
             now = monotonic()
             active_workers = []
-            for worker_id, worker_info in self._workers.items():
-                if self._worker_ttls.get(worker_id, 0) > now:
-                    active_workers.append(worker_info)
+            active_workers.extend(
+                worker_info
+                for worker_id, worker_info in self._workers.items()
+                if self._worker_ttls.get(worker_id, 0) > now
+            )
             return active_workers
 
     async def add_job_to_watch(self, job_id: str, timeout_at: float) -> None:
@@ -226,6 +229,7 @@ class MemoryStorage(StorageBackend):
             self._quotas.clear()
             self._generic_keys.clear()
             self._generic_key_ttls.clear()
+            self._locks.clear()
 
     async def get_job_queue_length(self) -> int:
         # No lock needed for asyncio.Queue.qsize()
@@ -234,13 +238,9 @@ class MemoryStorage(StorageBackend):
     async def get_active_worker_count(self) -> int:
         async with self._lock:
             now = monotonic()
-            count = 0
             # Create a copy of keys to avoid issues with concurrent modifications
             worker_ids = list(self._workers.keys())
-            for worker_id in worker_ids:
-                if self._worker_ttls.get(worker_id, 0) > now:
-                    count += 1
-            return count
+            return sum(self._worker_ttls.get(worker_id, 0) > now for worker_id in worker_ids)
 
     async def get_worker_info(self, worker_id: str) -> Optional[Dict[str, Any]]:
         async with self._lock:
@@ -273,3 +273,29 @@ class MemoryStorage(StorageBackend):
             "average_bid": 0,
             "error": "Statistics are not supported for MemoryStorage backend.",
         }
+
+    async def acquire_lock(self, key: str, holder_id: str, ttl: int) -> bool:
+        async with self._lock:
+            now = monotonic()
+            current_lock = self._locks.get(key)
+
+            # If lock exists and hasn't expired
+            if current_lock and current_lock[1] > now:
+                # If explicitly owned by us, we can extend/re-enter (optional behavior)
+                # But for strict locking, if it's held, return False (unless it's us? let's simpler: just False if held)
+                return False
+
+            # Acquire lock
+            self._locks[key] = (holder_id, now + ttl)
+            return True
+
+    async def release_lock(self, key: str, holder_id: str) -> bool:
+        async with self._lock:
+            current_lock = self._locks.get(key)
+            if current_lock:
+                owner, expiry = current_lock
+                # Only release if we are the owner
+                if owner == holder_id:
+                    del self._locks[key]
+                    return True
+            return False