avtomatika-worker 1.0b2__tar.gz → 1.0b4__tar.gz

{avtomatika_worker-1.0b2 → avtomatika_worker-1.0b4}/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-1.0b2 → avtomatika_worker-1.0b4}/PKG-INFO

@@ -1,19 +1,24 @@
 Metadata-Version: 2.4
 Name: avtomatika-worker
-Version: 1.0b2
+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: aioboto3~=15.5.0
+Requires-Dist: obstore>=0.1
 Requires-Dist: aiofiles~=25.1.0
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
@@ -21,13 +26,18 @@ Requires-Dist: pytest-asyncio; extra == "test"
 Requires-Dist: aioresponses; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
 Requires-Dist: pydantic; extra == "test"
+Requires-Dist: types-aiofiles; extra == "test"
 Provides-Extra: pydantic
 Requires-Dist: pydantic; extra == "pydantic"
 Dynamic: license-file
 
 # Avtomatika Worker SDK
 
-This is an SDK for creating workers compatible with the **Avtomatika** orchestrator. The SDK handles all the complexity of interacting with the orchestrator, allowing you to 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
 
@@ -285,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:
@@ -293,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:
@@ -472,7 +494,7 @@ async def generate_report(params: dict, files: TaskFiles, **kwargs):
 
 ### 6. Handling Large Files (S3 Payload Offloading)
 
-The SDK supports working with large files "out of the box" via S3-compatible storage.
+The SDK supports working with large files "out of the box" via S3-compatible storage, using the high-performance **`obstore`** library (Rust-based).
 
 -   **Automatic Download**: If a value in `params` is a URI of the form `s3://...`, the SDK will automatically download the file to the local disk and replace the URI in `params` with the local path. **If the URI ends with `/` (e.g., `s3://bucket/data/`), the SDK treats it as a folder prefix and recursively downloads all matching objects into a local directory.**
 -   **Automatic Upload**: If your handler returns a local file path in `data` (located within the `TASK_FILES_DIR` directory), the SDK will automatically upload this file to S3 and replace the path with an `s3://` URI in the final result. **If the path is a directory, the SDK recursively uploads all files within it.**
@@ -520,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
@@ -576,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`                                   |
@@ -600,11 +667,13 @@ The worker is fully configured via environment variables.
 | `S3_ACCESS_KEY`               | The access key for S3.                                                                                  | -                                      |
 | `S3_SECRET_KEY`               | The secret key for S3.                                                                                  | -                                      |
 | `S3_DEFAULT_BUCKET`           | The default bucket name for uploading results.                                                          | `avtomatika-payloads`                  |
+| `S3_REGION`                   | The region for S3 storage (required by some providers).                                                 | `us-east-1`                            |
 
 ## 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.0b2/src/avtomatika_worker.egg-info/PKG-INFO → avtomatika_worker-1.0b4/README.md

@@ -1,33 +1,10 @@
-Metadata-Version: 2.4
-Name: avtomatika-worker
-Version: 1.0b2
-Summary: Worker SDK for the Avtomatika orchestrator.
-Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika-worker
-Project-URL: Bug Tracker, https://github.com/avtomatika-ai/avtomatika-worker/issues
-Classifier: Development Status :: 4 - Beta
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: aiohttp~=3.13.2
-Requires-Dist: python-json-logger~=4.0.0
-Requires-Dist: aioboto3~=15.5.0
-Requires-Dist: aiofiles~=25.1.0
-Provides-Extra: test
-Requires-Dist: pytest; extra == "test"
-Requires-Dist: pytest-asyncio; extra == "test"
-Requires-Dist: aioresponses; extra == "test"
-Requires-Dist: pytest-mock; extra == "test"
-Requires-Dist: pydantic; extra == "test"
-Provides-Extra: pydantic
-Requires-Dist: pydantic; extra == "pydantic"
-Dynamic: license-file
-
 # Avtomatika Worker SDK
 
-This is an SDK for creating workers compatible with the **Avtomatika** orchestrator. The SDK handles all the complexity of interacting with the orchestrator, allowing you to 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
 
@@ -285,7 +262,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:
@@ -293,6 +270,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:
@@ -472,7 +461,7 @@ async def generate_report(params: dict, files: TaskFiles, **kwargs):
 
 ### 6. Handling Large Files (S3 Payload Offloading)
 
-The SDK supports working with large files "out of the box" via S3-compatible storage.
+The SDK supports working with large files "out of the box" via S3-compatible storage, using the high-performance **`obstore`** library (Rust-based).
 
 -   **Automatic Download**: If a value in `params` is a URI of the form `s3://...`, the SDK will automatically download the file to the local disk and replace the URI in `params` with the local path. **If the URI ends with `/` (e.g., `s3://bucket/data/`), the SDK treats it as a folder prefix and recursively downloads all matching objects into a local directory.**
 -   **Automatic Upload**: If your handler returns a local file path in `data` (located within the `TASK_FILES_DIR` directory), the SDK will automatically upload this file to S3 and replace the path with an `s3://` URI in the final result. **If the path is a directory, the SDK recursively uploads all files within it.**
@@ -520,6 +509,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
@@ -576,8 +607,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`                                   |
@@ -600,11 +634,13 @@ The worker is fully configured via environment variables.
 | `S3_ACCESS_KEY`               | The access key for S3.                                                                                  | -                                      |
 | `S3_SECRET_KEY`               | The secret key for S3.                                                                                  | -                                      |
 | `S3_DEFAULT_BUCKET`           | The default bucket name for uploading results.                                                          | `avtomatika-payloads`                  |
+| `S3_REGION`                   | The region for S3 storage (required by some providers).                                                 | `us-east-1`                            |
 
 ## 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.0b2 → avtomatika_worker-1.0b4}/pyproject.toml

@@ -4,20 +4,27 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "avtomatika-worker"
-version = "1.0.b2"
+version = "1.0b4"
 description = "Worker SDK for the Avtomatika orchestrator."
 readme = "README.md"
 requires-python = ">=3.11"
+authors = [
+    {name = "Dmitrii Gagarin", email = "madgagarin@gmail.com"},
+]
+keywords = ["worker", "sdk", "orchestrator", "distributed", "task-queue", "rxon", "hln"]
 classifiers = [
     "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
     "Programming Language :: Python :: 3",
     "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
+    "Typing :: Typed",
 ]
 dependencies = [
+    "rxon==1.0b2",
     "aiohttp~=3.13.2",
     "python-json-logger~=4.0.0",
-    "aioboto3~=15.5.0",
+    "obstore>=0.1",
     "aiofiles~=25.1.0",
 ]
 
@@ -28,6 +35,7 @@ test = [
     "aioresponses",
     "pytest-mock",
     "pydantic",
+    "types-aiofiles",
 ]
 pydantic = ["pydantic"]
 
@@ -38,6 +46,9 @@ pydantic = ["pydantic"]
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[tool.setuptools.package-data]
+"avtomatika_worker" = ["py.typed"]
+
 [tool.pytest.ini_options]
 markers = [
     "e2e: marks tests as end-to-end tests",

{avtomatika_worker-1.0b2 → avtomatika_worker-1.0b4}/src/avtomatika_worker/__init__.py

@@ -1,4 +1,4 @@
-"""A Python SDK for creating workers for the Py-Orchestrator."""
+"""A Python SDK for creating workers for the Avtomatika Orchestrator."""
 
 from importlib.metadata import PackageNotFoundError, version
 

{avtomatika_worker-1.0b2 → avtomatika_worker-1.0b4}/src/avtomatika_worker/config.py

@@ -4,13 +4,15 @@ from os import getenv
 from typing import Any
 from uuid import uuid4
 
+from rxon.validators import validate_identifier
+
 
 class WorkerConfig:
     """A class for centralized management of worker configuration.
     Reads parameters from environment variables and provides default values.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         # --- Basic worker information ---
         self.WORKER_ID: str = getenv("WORKER_ID", f"worker-{uuid4()}")
         self.WORKER_TYPE: str = getenv("WORKER_TYPE", "generic-cpu-worker")
@@ -29,6 +31,9 @@ class WorkerConfig:
             "WORKER_INDIVIDUAL_TOKEN",
             getenv("WORKER_TOKEN", "your-secret-worker-token"),
         )
+        self.TLS_CA_PATH: str | None = getenv("TLS_CA_PATH")
+        self.TLS_CERT_PATH: str | None = getenv("TLS_CERT_PATH")
+        self.TLS_KEY_PATH: str | None = getenv("TLS_KEY_PATH")
 
         # --- Resources and performance ---
         self.COST_PER_SKILL: dict[str, float] = self._load_json_from_env("COST_PER_SKILL", default={})
@@ -54,6 +59,7 @@ class WorkerConfig:
         self.S3_ACCESS_KEY: str | None = getenv("S3_ACCESS_KEY")
         self.S3_SECRET_KEY: str | None = getenv("S3_SECRET_KEY")
         self.S3_DEFAULT_BUCKET: str = getenv("S3_DEFAULT_BUCKET", "avtomatika-payloads")
+        self.S3_REGION: str = getenv("S3_REGION", "us-east-1")
 
         # --- Tuning parameters ---
         self.HEARTBEAT_INTERVAL: float = float(getenv("HEARTBEAT_INTERVAL", "15"))
@@ -70,6 +76,19 @@ class WorkerConfig:
         self.ENABLE_WEBSOCKETS: bool = getenv("WORKER_ENABLE_WEBSOCKETS", "false").lower() == "true"
         self.MULTI_ORCHESTRATOR_MODE: str = getenv("MULTI_ORCHESTRATOR_MODE", "FAILOVER")
 
+    def validate(self) -> None:
+        """Validates critical configuration parameters."""
+        validate_identifier(self.WORKER_ID, "WORKER_ID")
+        if self.WORKER_TOKEN == "your-secret-worker-token":
+            print("Warning: WORKER_TOKEN is set to the default value. Tasks might fail authentication.")
+
+        if not self.ORCHESTRATORS:
+            raise ValueError("No orchestrators configured.")
+
+        for o in self.ORCHESTRATORS:
+            if not o.get("url"):
+                raise ValueError("Orchestrator configuration missing URL.")
+
     def _get_orchestrators_config(self) -> list[dict[str, Any]]:
         """
         Loads orchestrator configuration from the ORCHESTRATORS_CONFIG environment variable.