avtomatika-worker 1.0b3__py3-none-any.whl → 1.0b4__py3-none-any.whl

{avtomatika_worker-1.0b3.dist-info → avtomatika_worker-1.0b4.dist-info}/METADATA

@@ -1,16 +1,21 @@
 Metadata-Version: 2.4
 Name: avtomatika-worker
-Version: 1.0b3
+Version: 1.0b4
 Summary: Worker SDK for the Avtomatika orchestrator.
+Author-email: Dmitrii Gagarin <madgagarin@gmail.com>
 Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika-worker
 Project-URL: Bug Tracker, https://github.com/avtomatika-ai/avtomatika-worker/issues
+Keywords: worker,sdk,orchestrator,distributed,task-queue,rxon,hln
 Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
+Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: rxon==1.0b2
 Requires-Dist: aiohttp~=3.13.2
 Requires-Dist: python-json-logger~=4.0.0
 Requires-Dist: obstore>=0.1
@@ -28,7 +33,11 @@ Dynamic: license-file
 
 # Avtomatika Worker SDK
 
-This is the official SDK for creating workers compatible with the **[Avtomatika Orchestrator](https://github.com/avtomatika-ai/avtomatika)**. It implements the **[RCA Protocol](https://github.com/avtomatika-ai/rca)**, handling all communication complexity (polling, heartbeats, S3 offloading) so you can focus on writing your business logic.
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/release/python-3110/)
+[![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)
+
+This is the official SDK for creating workers compatible with the **[Avtomatika Orchestrator](https://github.com/avtomatika-ai/avtomatika)**. It is built upon the **[Avtomatika Protocol](https://github.com/avtomatika-ai/rxon)** and implements the **[HLN Protocol](https://github.com/avtomatika-ai/hln)**, handling all communication complexity (polling, heartbeats, S3 offloading) so you can focus on writing your business logic.
 
 ## Installation
 
@@ -286,7 +295,7 @@ async def image_resizer(params: ResizeParams, **kwargs):
 
 ### 1. Task Handlers
 
-Each handler is an asynchronous function that accepts two arguments:
+Each handler is a function (either `async def` or `def`) that accepts two arguments:
 
 -   `params` (`dict`, `dataclass`, or `pydantic.BaseModel`): The parameters for the task, automatically validated and instantiated based on your type hint.
 -   `**kwargs`: Additional metadata about the task, including:
@@ -294,6 +303,18 @@ Each handler is an asynchronous function that accepts two arguments:
     -   `job_id` (`str`): The ID of the parent `Job` to which the task belongs.
     -   `priority` (`int`): The execution priority of the task.
 
+**Synchronous Handlers:**
+If you define your handler as a standard synchronous function (`def handler(...)`), the SDK will automatically execute it in a separate thread using `asyncio.to_thread`. This ensures that CPU-intensive operations (like model inference) do not block the worker's main event loop, allowing heartbeats and other background tasks to continue running smoothly.
+
+```python
+@worker.task("cpu_heavy_task")
+def heavy_computation(params: dict, **kwargs):
+    # This will run in a thread, not blocking the loop
+    import time
+    time.sleep(10) 
+    return {"status": "success"}
+```
+
 ### 2. Concurrency Limiting
 
 The worker allows you to control how many tasks are executed in parallel. This can be configured at two levels:
@@ -521,6 +542,48 @@ This only requires configuring environment variables for S3 access (see Full Con
 
 ### 7. WebSocket Support
 
+For real-time communication (e.g., immediate task cancellation), the worker supports WebSocket connections. This is enabled by setting `WORKER_ENABLE_WEBSOCKETS=true`. When connected, the orchestrator can push commands like `cancel_task` directly to the worker.
+
+### 8. Middleware
+
+The worker supports a middleware system, allowing you to wrap task executions with custom logic. This is particularly useful for resource management (e.g., acquiring GPU locks), logging, error handling, or **Dependency Injection**.
+
+Middleware functions wrap the execution of the task handler (and any subsequent middlewares). They receive a context dictionary and the next handler in the chain.
+
+The `context` dictionary contains:
+- `task_id`, `job_id`, `task_name`: Metadata.
+- `params`: The validated parameters object.
+- `handler_kwargs`: A dictionary of arguments that will be passed to the handler. **Middleware can modify this dictionary to inject dependencies.**
+
+**Example: GPU Resource Manager & Dependency Injection**
+
+```python
+async def gpu_lock_middleware(context: dict, next_handler: callable):
+    # Pre-processing: Acquire resource
+    print(f"Acquiring GPU for task {context['task_id']}...")
+    model_path = await resource_manager.allocate()
+    
+    # Inject the model path into the handler's arguments
+    context["handler_kwargs"]["model_path"] = model_path
+    
+    try:
+        # Execute the next handler in the chain
+        result = await next_handler()
+        return result
+    finally:
+        # Post-processing: Release resource
+        print(f"Releasing GPU for task {context['task_id']}...")
+        resource_manager.release()
+
+# Register the middleware
+worker.add_middleware(gpu_lock_middleware)
+
+# Handler now receives 'model_path' automatically
+@worker.task("generate")
+def generate(params, model_path, **kwargs):
+    print(f"Using model at: {model_path}")
+```
+
 ## Advanced Features
 
 ### Reporting Skill & Model Dependencies
@@ -577,8 +640,11 @@ The worker is fully configured via environment variables.
 | `WORKER_TYPE`                 | A string identifying the type of the worker.                                                            | `generic-cpu-worker`                   |
 | `WORKER_PORT`                 | The port for the worker's health check server.                                                          | `8083`                                 |
 | `WORKER_TOKEN`                | A common authentication token used to connect to orchestrators.                                         | `your-secret-worker-token`             |
-| `WORKER_INDIVIDUAL_TOKEN`     | An individual token for this worker, which overrides `WORKER_TOKEN` if set.                               | -                                      |
-| `ORCHESTRATOR_URL`            | The URL of a single orchestrator (used if `ORCHESTRATORS_CONFIG` is not set).                             | `http://localhost:8080`                |
+-   **`WORKER_INDIVIDUAL_TOKEN`**: An individual token for this worker, which overrides `WORKER_TOKEN` if set.
+-   **`TLS_CA_PATH`**: Path to the CA certificate to verify the orchestrator.
+-   **`TLS_CERT_PATH`**: Path to the client certificate for mTLS.
+-   **`TLS_KEY_PATH`**: Path to the client private key for mTLS.
+-   **`ORCHESTRATOR_URL`**: The address of the Avtomatika orchestrator.
 | `ORCHESTRATORS_CONFIG`        | A JSON string with a list of orchestrators for multi-orchestrator modes.                                | `[]`                                   |
 | `MULTI_ORCHESTRATOR_MODE`     | The mode for handling multiple orchestrators. Possible values: `FAILOVER`, `ROUND_ROBIN`.                  | `FAILOVER`                             |
 | `MAX_CONCURRENT_TASKS`        | The maximum number of tasks the worker can execute simultaneously.                                      | `10`                                   |
@@ -605,8 +671,9 @@ The worker is fully configured via environment variables.
 
 ## Development
 
-To install the necessary dependencies for running tests, use the following command:
+To install the necessary dependencies for running tests (assuming you are in the package root):
 
-```bash
-pip install .[test]
-```
+1.  Install the worker in editable mode with test dependencies:
+    ```bash
+    pip install -e .[test]
+    ```

avtomatika_worker-1.0b4.dist-info/RECORD

@@ -0,0 +1,12 @@
+avtomatika_worker/__init__.py,sha256=YYxuVc3EUabivZe6ZXNyFDUZH97qTKIQwxKZxV4klDc,342
+avtomatika_worker/config.py,sha256=wXmtQTABhHEl0vwxQdWgSfmCFsKNEH-lF5CF70VV2SU,6114
+avtomatika_worker/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+avtomatika_worker/s3.py,sha256=v_xSMVsPRgjkhQZWJreGjMSYCIfH_n7BHo5GMqmPw5k,10371
+avtomatika_worker/task_files.py,sha256=9WmvVt4iZP65CA3lTttOmVnW8tlvfQMwutfTla4yJyY,5858
+avtomatika_worker/types.py,sha256=cHzpvdJahAV6RlirlWcQ0xhxCrwNlrCalbsUL49t_ko,581
+avtomatika_worker/worker.py,sha256=oMuMiPGKUYKwCq18T11zQeT36OBSGfxPNd5unVEEWsk,28250
+avtomatika_worker-1.0b4.dist-info/licenses/LICENSE,sha256=J19fUi8XywciRuLLWHvcPQGc0Uo-9K1a0dKaIbzw_Yg,1092
+avtomatika_worker-1.0b4.dist-info/METADATA,sha256=DZHKg_Yxl0dTraVoLfj29lOPvVzAJD8taF8VjxxLlUU,32868
+avtomatika_worker-1.0b4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+avtomatika_worker-1.0b4.dist-info/top_level.txt,sha256=d3b5BUeUrHM1Cn-cbStz-hpucikEBlPOvtcmQ_j3qAs,18
+avtomatika_worker-1.0b4.dist-info/RECORD,,

{avtomatika_worker-1.0b3.dist-info → avtomatika_worker-1.0b4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{avtomatika_worker-1.0b3.dist-info → avtomatika_worker-1.0b4.dist-info}/licenses/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2025 Dmitrii Gagarin
+Copyright (c) 2025-2026 Dmitrii Gagarin aka madgagarin
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

avtomatika_worker/client.py

@@ -1,93 +0,0 @@
-from asyncio import sleep
-from logging import getLogger
-from typing import Any
-
-from aiohttp import ClientError, ClientSession, ClientTimeout, ClientWebSocketResponse
-
-from .constants import AUTH_HEADER_WORKER
-
-logger = getLogger(__name__)
-
-
-class OrchestratorClient:
-    """
-    Dedicated client for communicating with a single Avtomatika Orchestrator instance.
-    Handles HTTP requests, retries, and authentication.
-    """
-
-    def __init__(self, session: ClientSession, base_url: str, worker_id: str, token: str):
-        self.session = session
-        self.base_url = base_url.rstrip("/")
-        self.worker_id = worker_id
-        self.token = token
-        self._headers = {AUTH_HEADER_WORKER: self.token}
-
-    async def register(self, payload: dict[str, Any]) -> bool:
-        """Registers the worker with the orchestrator."""
-        url = f"{self.base_url}/_worker/workers/register"
-        try:
-            async with self.session.post(url, json=payload, headers=self._headers) as resp:
-                if resp.status >= 400:
-                    logger.error(f"Error registering with {self.base_url}: {resp.status}")
-                    return False
-                return True
-        except ClientError as e:
-            logger.error(f"Error registering with orchestrator {self.base_url}: {e}")
-            return False
-
-    async def poll_task(self, timeout: float) -> dict[str, Any] | None:
-        """Polls for the next available task."""
-        url = f"{self.base_url}/_worker/workers/{self.worker_id}/tasks/next"
-        client_timeout = ClientTimeout(total=timeout + 5)
-        try:
-            async with self.session.get(url, headers=self._headers, timeout=client_timeout) as resp:
-                if resp.status == 200:
-                    return await resp.json()
-                elif resp.status != 204:
-                    logger.warning(f"Unexpected status from {self.base_url} during poll: {resp.status}")
-        except ClientError as e:
-            logger.error(f"Error polling for tasks from {self.base_url}: {e}")
-        except Exception as e:
-            logger.exception(f"Unexpected error polling from {self.base_url}: {e}")
-        return None
-
-    async def send_heartbeat(self, payload: dict[str, Any]) -> bool:
-        """Sends a heartbeat message to update worker state."""
-        url = f"{self.base_url}/_worker/workers/{self.worker_id}"
-        try:
-            async with self.session.patch(url, json=payload, headers=self._headers) as resp:
-                if resp.status >= 400:
-                    logger.warning(f"Heartbeat to {self.base_url} failed with status: {resp.status}")
-                    return False
-                return True
-        except ClientError as e:
-            logger.error(f"Error sending heartbeat to orchestrator {self.base_url}: {e}")
-            return False
-
-    async def send_result(self, payload: dict[str, Any], max_retries: int, initial_delay: float) -> bool:
-        """Sends task result with retries and exponential backoff."""
-        url = f"{self.base_url}/_worker/tasks/result"
-        delay = initial_delay
-        for i in range(max_retries):
-            try:
-                async with self.session.post(url, json=payload, headers=self._headers) as resp:
-                    if resp.status == 200:
-                        return True
-                    logger.error(f"Error sending result to {self.base_url}: {resp.status}")
-            except ClientError as e:
-                logger.error(f"Error sending result to {self.base_url}: {e}")
-
-            if i < max_retries - 1:
-                await sleep(delay * (2**i))
-        return False
-
-    async def connect_websocket(self) -> ClientWebSocketResponse | None:
-        """Establishes a WebSocket connection for real-time commands."""
-        ws_url = self.base_url.replace("http", "ws", 1) + "/_worker/ws"
-        try:
-            ws = await self.session.ws_connect(ws_url, headers=self._headers)
-            logger.info(f"WebSocket connection established to {ws_url}")
-            return ws
-        except Exception as e:
-            logger.warning(f"WebSocket connection to {ws_url} failed: {e}")
-            return None

avtomatika_worker/constants.py

@@ -1,22 +0,0 @@
-"""
-Centralized constants for the Avtomatika protocol (Worker SDK).
-These should match the constants in the core `avtomatika` package.
-"""
-
-# --- Auth Headers ---
-AUTH_HEADER_CLIENT = "X-Avtomatika-Token"
-AUTH_HEADER_WORKER = "X-Worker-Token"
-
-# --- Error Codes ---
-ERROR_CODE_TRANSIENT = "TRANSIENT_ERROR"
-ERROR_CODE_PERMANENT = "PERMANENT_ERROR"
-ERROR_CODE_INVALID_INPUT = "INVALID_INPUT_ERROR"
-
-# --- Task Statuses ---
-TASK_STATUS_SUCCESS = "success"
-TASK_STATUS_FAILURE = "failure"
-TASK_STATUS_CANCELLED = "cancelled"
-TASK_STATUS_NEEDS_REVIEW = "needs_review"  # Example of a common custom status
-
-# --- Commands (WebSocket) ---
-COMMAND_CANCEL_TASK = "cancel_task"

avtomatika_worker-1.0b3.dist-info/RECORD

@@ -1,14 +0,0 @@
-avtomatika_worker/__init__.py,sha256=y_s5KlsgFu7guemZfjLVQ3Jzq7DyLG168-maVGwWRC4,334
-avtomatika_worker/client.py,sha256=mkvwrMY8tAaZN_lwMSxWHmAoWsDemD-WiKSeH5fM6GI,4173
-avtomatika_worker/config.py,sha256=NaAhufpwyG6CsHW-cXmqR3MfGp_5SdDZ_vEhmmV8G3g,5819
-avtomatika_worker/constants.py,sha256=DfGR_YkW9rbioCorKpNGfZ0i_0iGgMq2swyJhVl9nNA,669
-avtomatika_worker/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-avtomatika_worker/s3.py,sha256=sAuXBp__XTVhyNwK9gsxy1Jm_udJM6ypssGTZ00pa6U,8847
-avtomatika_worker/task_files.py,sha256=ucjBuI78UmtMvfucTzDTNJ1g0KJaRIwyshRNTipIZSU,3351
-avtomatika_worker/types.py,sha256=dSNsHgqV6hZhOt4eUK2PDWB6lrrwCA5_T_iIBI_wTZ0,442
-avtomatika_worker/worker.py,sha256=XSRfLO-W0J6WG128Iu-rL_w3-PqmsWQMUElVLi3Z1gk,21904
-avtomatika_worker-1.0b3.dist-info/licenses/LICENSE,sha256=tqCjw9Y1vbU-hLcWi__7wQstLbt2T1XWPdbQYqCxuWY,1072
-avtomatika_worker-1.0b3.dist-info/METADATA,sha256=yiEtJuMv5WHYHfScna7cF5QAvAUhMCJdUDENHvrMRFY,29601
-avtomatika_worker-1.0b3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-avtomatika_worker-1.0b3.dist-info/top_level.txt,sha256=d3b5BUeUrHM1Cn-cbStz-hpucikEBlPOvtcmQ_j3qAs,18
-avtomatika_worker-1.0b3.dist-info/RECORD,,