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

avtomatika_worker/client.py

@@ -0,0 +1,93 @@
+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/config.py

@@ -10,7 +10,7 @@ class WorkerConfig:
     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")
@@ -54,6 +54,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 +71,18 @@ 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."""
+        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.

avtomatika_worker/constants.py

@@ -0,0 +1,22 @@
+"""
+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/s3.py

@@ -1,36 +1,61 @@
-from asyncio import gather, to_thread
+from asyncio import Semaphore, gather, to_thread
+from logging import getLogger
 from os import walk
 from os.path import basename, dirname, join, relpath
 from shutil import rmtree
-from typing import Any
+from typing import Any, cast
 from urllib.parse import urlparse
 
-from aioboto3 import Session
+import obstore
+from aiofiles import open as aio_open
 from aiofiles.os import makedirs
 from aiofiles.ospath import exists, isdir
-from botocore.client import Config
+from obstore.store import S3Store
 
 from .config import WorkerConfig
 
+logger = getLogger(__name__)
+
+# Limit concurrent S3 operations to avoid "Too many open files"
+MAX_S3_CONCURRENCY = 50
+
 
 class S3Manager:
-    """Handles S3 payload offloading."""
+    """Handles S3 payload offloading using obstore (high-performance async S3 client)."""
 
     def __init__(self, config: WorkerConfig):
         self._config = config
-        self._session = Session()
+        self._stores: dict[str, S3Store] = {}
+        self._semaphore = Semaphore(MAX_S3_CONCURRENCY)
+
+    def _get_store(self, bucket_name: str) -> S3Store:
+        """Creates or returns a cached S3Store for a specific bucket."""
+        if bucket_name in self._stores:
+            return self._stores[bucket_name]
 
-    def _get_client_args(self) -> dict[str, Any]:
-        """Returns standard arguments for S3 client creation."""
-        return {
-            "service_name": "s3",
-            "endpoint_url": self._config.S3_ENDPOINT_URL,
+        config_kwargs = {
             "aws_access_key_id": self._config.S3_ACCESS_KEY,
             "aws_secret_access_key": self._config.S3_SECRET_KEY,
-            "config": Config(signature_version="s3v4"),
+            "region": "us-east-1",  # Default region if not specified, required by some clients
         }
 
-    async def cleanup(self, task_id: str):
+        if self._config.S3_ENDPOINT_URL:
+            config_kwargs["endpoint"] = self._config.S3_ENDPOINT_URL
+            if self._config.S3_ENDPOINT_URL.startswith("http://"):
+                config_kwargs["allow_http"] = "true"
+
+        # Filter out None values
+        config_kwargs = {k: v for k, v in config_kwargs.items() if v is not None}
+
+        try:
+            store = S3Store(bucket_name, **config_kwargs)
+            self._stores[bucket_name] = store
+            return store
+        except Exception as e:
+            logger.error(f"Failed to create S3Store for bucket {bucket_name}: {e}")
+            raise
+
+    async def cleanup(self, task_id: str) -> None:
         """Removes the task-specific payload directory."""
         task_dir = join(self._config.TASK_FILES_DIR, task_id)
         if await exists(task_dir):
@@ -38,54 +63,83 @@ class S3Manager:
 
     async def _process_s3_uri(self, uri: str, task_id: str) -> str:
         """Downloads a file or a folder (if uri ends with /) from S3 and returns the local path."""
-        parsed_url = urlparse(uri)
-        bucket_name = parsed_url.netloc
-        object_key = parsed_url.path.lstrip("/")
+        try:
+            parsed_url = urlparse(uri)
+            bucket_name = parsed_url.netloc
+            object_key = parsed_url.path.lstrip("/")
+            store = self._get_store(bucket_name)
+
+            # Use task-specific directory for isolation
+            local_dir_root = join(self._config.TASK_FILES_DIR, task_id)
+            await makedirs(local_dir_root, exist_ok=True)
 
-        # Use task-specific directory for isolation
-        local_dir_root = join(self._config.TASK_FILES_DIR, task_id)
-        await makedirs(local_dir_root, exist_ok=True)
+            logger.info(f"Starting download from S3: {uri}")
 
-        async with self._session.client(**self._get_client_args()) as s3:
             # Handle folder download (prefix)
             if uri.endswith("/"):
                 folder_name = object_key.rstrip("/").split("/")[-1]
                 local_folder_path = join(local_dir_root, folder_name)
 
-                paginator = s3.get_paginator("list_objects_v2")
-                tasks = []
-                async for page in paginator.paginate(Bucket=bucket_name, Prefix=object_key):
-                    for obj in page.get("Contents", []):
-                        key = obj["Key"]
-                        if key.endswith("/"):
-                            continue
+                # List objects with prefix
+                # obstore.list returns an async iterator of ObjectMeta
+                files_to_download = []
+
+                # Note: obstore.list returns an async iterator.
+                async for obj in obstore.list(store, prefix=object_key):
+                    key = obj.key
+
+                    if key.endswith("/"):
+                        continue
 
-                        # Calculate relative path inside the folder
-                        rel_path = key[len(object_key) :]
-                        local_file_path = join(local_folder_path, rel_path)
+                    # Calculate relative path inside the folder
+                    rel_path = key[len(object_key) :]
+                    local_file_path = join(local_folder_path, rel_path)
 
-                        await makedirs(dirname(local_file_path), exist_ok=True)
-                        tasks.append(s3.download_file(bucket_name, key, local_file_path))
+                    await makedirs(dirname(local_file_path), exist_ok=True)
+                    files_to_download.append((key, local_file_path))
 
-                if tasks:
-                    await gather(*tasks)
+                async def _download_file(key: str, path: str) -> None:
+                    async with self._semaphore:
+                        result = await obstore.get(store, key)
+                        async with aio_open(path, "wb") as f:
+                            async for chunk in result.stream():
+                                await f.write(chunk)
+
+                # Execute downloads in parallel
+                if files_to_download:
+                    await gather(*[_download_file(k, p) for k, p in files_to_download])
+
+                logger.info(f"Successfully downloaded folder from S3: {uri} ({len(files_to_download)} files)")
                 return local_folder_path
 
             # Handle single file download
             local_path = join(local_dir_root, basename(object_key))
-            await s3.download_file(bucket_name, object_key, local_path)
+
+            result = await obstore.get(store, object_key)
+            async with aio_open(local_path, "wb") as f:
+                async for chunk in result.stream():
+                    await f.write(chunk)
+
+            logger.info(f"Successfully downloaded file from S3: {uri} -> {local_path}")
             return local_path
 
+        except Exception as e:
+            # Catching generic Exception because obstore might raise different errors.
+            logger.exception(f"Error during download of {uri}: {e}")
+            raise
+
     async def _upload_to_s3(self, local_path: str) -> str:
         """Uploads a file or a folder to S3 and returns the S3 URI."""
         bucket_name = self._config.S3_DEFAULT_BUCKET
+        store = self._get_store(bucket_name)
+
+        logger.info(f"Starting upload to S3 from local path: {local_path}")
 
-        async with self._session.client(**self._get_client_args()) as s3:
+        try:
             # Handle folder upload
             if await isdir(local_path):
                 folder_name = basename(local_path.rstrip("/"))
                 s3_prefix = f"{folder_name}/"
-                tasks = []
 
                 # Use to_thread to avoid blocking event loop during file walk
                 def _get_files_to_upload():
@@ -99,18 +153,33 @@ class S3Manager:
 
                 files_list = await to_thread(_get_files_to_upload)
 
-                for full_path, key in files_list:
-                    tasks.append(s3.upload_file(full_path, bucket_name, key))
+                async def _upload_file(path: str, key: str) -> None:
+                    async with self._semaphore:
+                        # obstore.put accepts bytes or file-like objects.
+                        # Since we are in async, reading small files is fine.
+                        with open(path, "rb") as f:
+                            await obstore.put(store, key, f)
 
-                if tasks:
-                    await gather(*tasks)
+                if files_list:
+                    # Upload in parallel
+                    await gather(*[_upload_file(f, k) for f, k in files_list])
 
-                return f"s3://{bucket_name}/{s3_prefix}"
+                s3_uri = f"s3://{bucket_name}/{s3_prefix}"
+                logger.info(f"Successfully uploaded folder to S3: {local_path} -> {s3_uri} ({len(files_list)} files)")
+                return s3_uri
 
             # Handle single file upload
             object_key = basename(local_path)
-            await s3.upload_file(local_path, bucket_name, object_key)
-            return f"s3://{bucket_name}/{object_key}"
+            with open(local_path, "rb") as f:
+                await obstore.put(store, object_key, f)
+
+            s3_uri = f"s3://{bucket_name}/{object_key}"
+            logger.info(f"Successfully uploaded file to S3: {local_path} -> {s3_uri}")
+            return s3_uri
+
+        except Exception as e:
+            logger.exception(f"Error during upload of {local_path}: {e}")
+            raise
 
     async def process_params(self, params: dict[str, Any], task_id: str) -> dict[str, Any]:
         """Recursively searches for S3 URIs in params and downloads the files."""
@@ -124,7 +193,7 @@ class S3Manager:
                 return {k: await _process(v) for k, v in item.items()}
             return [await _process(i) for i in item] if isinstance(item, list) else item
 
-        return await _process(params)
+        return cast(dict[str, Any], await _process(params))
 
     async def process_result(self, result: dict[str, Any]) -> dict[str, Any]:
         """Recursively searches for local file paths in the result and uploads them to S3."""
@@ -138,4 +207,4 @@ class S3Manager:
                 return {k: await _process(v) for k, v in item.items()}
             return [await _process(i) for i in item] if isinstance(item, list) else item
 
-        return await _process(result)
+        return cast(dict[str, Any], await _process(result))

avtomatika_worker/types.py

@@ -1,8 +1,21 @@
-# Error codes for worker task results
-TRANSIENT_ERROR = "TRANSIENT_ERROR"
-PERMANENT_ERROR = "PERMANENT_ERROR"
-INVALID_INPUT_ERROR = "INVALID_INPUT_ERROR"
+from .constants import (
+    ERROR_CODE_INVALID_INPUT as INVALID_INPUT_ERROR,
+)
+from .constants import (
+    ERROR_CODE_PERMANENT as PERMANENT_ERROR,
+)
+from .constants import (
+    ERROR_CODE_TRANSIENT as TRANSIENT_ERROR,
+)
 
 
 class ParamValidationError(Exception):
     """Custom exception for parameter validation errors."""
+
+
+__all__ = [
+    "INVALID_INPUT_ERROR",
+    "PERMANENT_ERROR",
+    "TRANSIENT_ERROR",
+    "ParamValidationError",
+]

avtomatika_worker/worker.py

@@ -1,5 +1,4 @@
 from asyncio import CancelledError, Event, Task, create_task, gather, run, sleep
-from asyncio import TimeoutError as AsyncTimeoutError
 from dataclasses import is_dataclass
 from inspect import Parameter, signature
 from json import JSONDecodeError
@@ -7,12 +6,21 @@ from logging import getLogger
 from os.path import join
 from typing import Any, Callable
 
-from aiohttp import ClientError, ClientSession, ClientTimeout, ClientWebSocketResponse, WSMsgType, web
+from aiohttp import ClientSession, ClientWebSocketResponse, WSMsgType, web
 
+from .client import OrchestratorClient
 from .config import WorkerConfig
+from .constants import (
+    COMMAND_CANCEL_TASK,
+    ERROR_CODE_INVALID_INPUT,
+    ERROR_CODE_PERMANENT,
+    ERROR_CODE_TRANSIENT,
+    TASK_STATUS_CANCELLED,
+    TASK_STATUS_FAILURE,
+)
 from .s3 import S3Manager
 from .task_files import TaskFiles
-from .types import INVALID_INPUT_ERROR, PERMANENT_ERROR, TRANSIENT_ERROR, ParamValidationError
+from .types import ParamValidationError
 
 try:
     from pydantic import BaseModel, ValidationError
@@ -45,7 +53,7 @@ class Worker:
         self._s3_manager = S3Manager(self._config)
         self._config.WORKER_TYPE = worker_type  # Allow overriding worker_type
         if max_concurrent_tasks is not None:
-            self._config.max_concurrent_tasks = max_concurrent_tasks
+            self._config.MAX_CONCURRENT_TASKS = max_concurrent_tasks
 
         self._task_type_limits = task_type_limits or {}
         self._task_handlers: dict[str, dict[str, Any]] = {}
@@ -59,10 +67,8 @@ class Worker:
         self._http_session = http_session
         self._session_is_managed_externally = http_session is not None
         self._ws_connection: ClientWebSocketResponse | None = None
-        # Removed: self._headers = {"X-Worker-Token": self._config.WORKER_TOKEN}
         self._shutdown_event = Event()
         self._registered_event = Event()
-        self._round_robin_index = 0
         self._debounce_task: Task | None = None
 
         # --- Weighted Round-Robin State ---
@@ -72,7 +78,28 @@ class Worker:
                 o["current_weight"] = 0
                 self._total_orchestrator_weight += o.get("weight", 1)
 
-    def _validate_config(self):
+        self._clients: list[tuple[dict[str, Any], OrchestratorClient]] = []
+        if self._http_session:
+            self._init_clients()
+
+    def _init_clients(self):
+        """Initializes OrchestratorClient instances for each configured orchestrator."""
+        if not self._http_session:
+            return
+        self._clients = [
+            (
+                o,
+                OrchestratorClient(
+                    session=self._http_session,
+                    base_url=o["url"],
+                    worker_id=self._config.WORKER_ID,
+                    token=o.get("token", self._config.WORKER_TOKEN),
+                ),
+            )
+            for o in self._config.ORCHESTRATORS
+        ]
+
+    def _validate_task_types(self):
         """Checks for unused task type limits and warns the user."""
         registered_task_types = {
             handler_data["type"] for handler_data in self._task_handlers.values() if handler_data["type"]
@@ -140,32 +167,31 @@ class Worker:
         status = "idle" if supported_tasks else "busy"
         return {"status": status, "supported_tasks": supported_tasks}
 
-    def _get_headers(self, orchestrator: dict[str, Any]) -> dict[str, str]:
-        """Builds authentication headers for a specific orchestrator."""
-        token = orchestrator.get("token", self._config.WORKER_TOKEN)
-        return {"X-Worker-Token": token}
-
-    def _get_next_orchestrator(self) -> dict[str, Any] | None:
+    def _get_next_client(self) -> OrchestratorClient | None:
         """
-        Selects the next orchestrator using a smooth weighted round-robin algorithm.
+        Selects the next orchestrator client using a smooth weighted round-robin algorithm.
         """
-        if not self._config.ORCHESTRATORS:
+        if not self._clients:
             return None
 
         # The orchestrator with the highest current_weight is selected.
-        selected_orchestrator = None
+        selected_client = None
         highest_weight = -1
 
-        for o in self._config.ORCHESTRATORS:
+        for o, client in self._clients:
             o["current_weight"] += o["weight"]
             if o["current_weight"] > highest_weight:
                 highest_weight = o["current_weight"]
-                selected_orchestrator = o
+                selected_client = client
 
-        if selected_orchestrator:
-            selected_orchestrator["current_weight"] -= self._total_orchestrator_weight
+        if selected_client:
+            # Find the config for the selected client to decrement its weight
+            for o, client in self._clients:
+                if client == selected_client:
+                    o["current_weight"] -= self._total_orchestrator_weight
+                    break
 
-        return selected_orchestrator
+        return selected_client
 
     async def _debounced_heartbeat_sender(self):
         """Waits for the debounce delay then sends a heartbeat."""
@@ -180,33 +206,27 @@ class Worker:
         # Schedule the new debounced call.
         self._debounce_task = create_task(self._debounced_heartbeat_sender())
 
-    async def _poll_for_tasks(self, orchestrator: dict[str, Any]):
+    async def _poll_for_tasks(self, client: OrchestratorClient):
         """Polls a specific Orchestrator for new tasks."""
-        url = f"{orchestrator['url']}/_worker/workers/{self._config.WORKER_ID}/tasks/next"
-        try:
-            if not self._http_session:
-                return
-            timeout = ClientTimeout(total=self._config.TASK_POLL_TIMEOUT + 5)
-            headers = self._get_headers(orchestrator)
-            async with self._http_session.get(url, headers=headers, timeout=timeout) as resp:
-                if resp.status == 200:
-                    task_data = await resp.json()
-                    task_data["orchestrator"] = orchestrator
-
-                    self._current_load += 1
-                    if (task_handler_info := self._task_handlers.get(task_data["type"])) and (
-                        task_type_for_limit := task_handler_info.get("type")
-                    ):
-                        self._current_load_by_type[task_type_for_limit] += 1
-                    self._schedule_heartbeat_debounce()
-
-                    task = create_task(self._process_task(task_data))
-                    self._active_tasks[task_data["task_id"]] = task
-                elif resp.status != 204:
-                    await sleep(self._config.TASK_POLL_ERROR_DELAY)
-        except (AsyncTimeoutError, ClientError) as e:
-            logger.error(f"Error polling for tasks: {e}")
-            await sleep(self._config.TASK_POLL_ERROR_DELAY)
+        task_data = await client.poll_task(timeout=self._config.TASK_POLL_TIMEOUT)
+        if task_data:
+            task_data["client"] = client
+
+            self._current_load += 1
+            if (task_handler_info := self._task_handlers.get(task_data["type"])) and (
+                task_type_for_limit := task_handler_info.get("type")
+            ):
+                self._current_load_by_type[task_type_for_limit] += 1
+            self._schedule_heartbeat_debounce()
+
+            task = create_task(self._process_task(task_data))
+            self._active_tasks[task_data["task_id"]] = task
+        else:
+            # If no task but it was a 204 or error, the client already handled/logged it.
+            # We might want a short sleep here if it was an error, but client.poll_task
+            # doesn't distinguish between 204 and error currently.
+            # However, the previous logic only slept on status != 204.
+            pass
 
     async def _start_polling(self):
         """The main loop for polling tasks."""
@@ -218,13 +238,13 @@ class Worker:
                 continue
 
             if self._config.MULTI_ORCHESTRATOR_MODE == "ROUND_ROBIN":
-                if orchestrator := self._get_next_orchestrator():
-                    await self._poll_for_tasks(orchestrator)
+                if client := self._get_next_client():
+                    await self._poll_for_tasks(client)
             else:
-                for orchestrator in self._config.ORCHESTRATORS:
+                for _, client in self._clients:
                     if self._get_current_state()["status"] == "busy":
                         break
-                    await self._poll_for_tasks(orchestrator)
+                    await self._poll_for_tasks(client)
 
             if self._current_load == 0:
                 await sleep(self._config.IDLE_POLL_DELAY)
@@ -289,7 +309,7 @@ class Worker:
     async def _process_task(self, task_data: dict[str, Any]):
         """Executes the task logic."""
         task_id, job_id, task_name = task_data["task_id"], task_data["job_id"], task_data["type"]
-        params, orchestrator = task_data.get("params", {}), task_data["orchestrator"]
+        params, client = task_data.get("params", {}), task_data["client"]
 
         result: dict[str, Any] = {}
         handler_data = self._task_handlers.get(task_name)
@@ -301,9 +321,11 @@ class Worker:
             if not handler_data:
                 message = f"Unsupported task: {task_name}"
                 logger.warning(message)
-                result = {"status": "failure", "error": {"code": PERMANENT_ERROR, "message": message}}
+                result = {"status": TASK_STATUS_FAILURE, "error": {"code": ERROR_CODE_PERMANENT, "message": message}}
                 payload = {"job_id": job_id, "task_id": task_id, "worker_id": self._config.WORKER_ID, "result": result}
-                await self._send_result(payload, orchestrator)
+                await client.send_result(
+                    payload, self._config.RESULT_MAX_RETRIES, self._config.RESULT_RETRY_INITIAL_DELAY
+                )
                 result_sent = True  # Mark result as sent
                 return
 
@@ -324,22 +346,24 @@ class Worker:
             result = await self._s3_manager.process_result(result)
         except ParamValidationError as e:
             logger.error(f"Task {task_id} failed validation: {e}")
-            result = {"status": "failure", "error": {"code": INVALID_INPUT_ERROR, "message": str(e)}}
+            result = {"status": TASK_STATUS_FAILURE, "error": {"code": ERROR_CODE_INVALID_INPUT, "message": str(e)}}
         except CancelledError:
             logger.info(f"Task {task_id} was cancelled.")
-            result = {"status": "cancelled"}
+            result = {"status": TASK_STATUS_CANCELLED}
             # We must re-raise the exception to be handled by the outer gather
             raise
         except Exception as e:
             logger.exception(f"An unexpected error occurred while processing task {task_id}:")
-            result = {"status": "failure", "error": {"code": TRANSIENT_ERROR, "message": str(e)}}
+            result = {"status": TASK_STATUS_FAILURE, "error": {"code": ERROR_CODE_TRANSIENT, "message": str(e)}}
         finally:
             # Cleanup task workspace
             await self._s3_manager.cleanup(task_id)
 
             if not result_sent:  # Only send if not already sent
                 payload = {"job_id": job_id, "task_id": task_id, "worker_id": self._config.WORKER_ID, "result": result}
-                await self._send_result(payload, orchestrator)
+                await client.send_result(
+                    payload, self._config.RESULT_MAX_RETRIES, self._config.RESULT_RETRY_INITIAL_DELAY
+                )
             self._active_tasks.pop(task_id, None)
 
             self._current_load -= 1
@@ -347,21 +371,6 @@ class Worker:
                 self._current_load_by_type[task_type_for_limit] -= 1
             self._schedule_heartbeat_debounce()
 
-    async def _send_result(self, payload: dict[str, Any], orchestrator: dict[str, Any]):
-        """Sends the result to a specific orchestrator."""
-        url = f"{orchestrator['url']}/_worker/tasks/result"
-        delay = self._config.RESULT_RETRY_INITIAL_DELAY
-        headers = self._get_headers(orchestrator)
-        for i in range(self._config.RESULT_MAX_RETRIES):
-            try:
-                if self._http_session and not self._http_session.closed:
-                    async with self._http_session.post(url, json=payload, headers=headers) as resp:
-                        if resp.status == 200:
-                            return
-            except ClientError as e:
-                logger.error(f"Error sending result: {e}")
-            await sleep(delay * (2**i))
-
     async def _manage_orchestrator_communications(self):
         """Registers the worker and sends heartbeats."""
         await self._register_with_all_orchestrators()
@@ -388,17 +397,7 @@ class Worker:
             "ip_address": self._config.IP_ADDRESS,
             "resources": self._config.RESOURCES,
         }
-        for orchestrator in self._config.ORCHESTRATORS:
-            url = f"{orchestrator['url']}/_worker/workers/register"
-            try:
-                if self._http_session:
-                    async with self._http_session.post(
-                        url, json=payload, headers=self._get_headers(orchestrator)
-                    ) as resp:
-                        if resp.status >= 400:
-                            logger.error(f"Error registering with {orchestrator['url']}: {resp.status}")
-            except ClientError as e:
-                logger.error(f"Error registering with orchestrator {orchestrator['url']}: {e}")
+        await gather(*[client.register(payload) for _, client in self._clients])
 
     async def _send_heartbeats_to_all(self):
         """Sends heartbeat messages to all orchestrators."""
@@ -418,24 +417,15 @@ class Worker:
             if hot_skills:
                 payload["hot_skills"] = hot_skills
 
-        async def _send_single(orchestrator: dict[str, Any]):
-            url = f"{orchestrator['url']}/_worker/workers/{self._config.WORKER_ID}"
-            headers = self._get_headers(orchestrator)
-            try:
-                if self._http_session and not self._http_session.closed:
-                    async with self._http_session.patch(url, json=payload, headers=headers) as resp:
-                        if resp.status >= 400:
-                            logger.warning(f"Heartbeat to {orchestrator['url']} failed with status: {resp.status}")
-            except ClientError as e:
-                logger.error(f"Error sending heartbeat to orchestrator {orchestrator['url']}: {e}")
-
-        await gather(*[_send_single(o) for o in self._config.ORCHESTRATORS])
+        await gather(*[client.send_heartbeat(payload) for _, client in self._clients])
 
     async def main(self):
         """The main asynchronous function."""
-        self._validate_config()  # Validate config now that all tasks are registered
+        self._config.validate()
+        self._validate_task_types()  # Validate config now that all tasks are registered
         if not self._http_session:
             self._http_session = ClientSession()
+            self._init_clients()
 
         comm_task = create_task(self._manage_orchestrator_communications())
 
@@ -482,25 +472,20 @@ class Worker:
         except KeyboardInterrupt:
             self._shutdown_event.set()
 
-    # WebSocket methods omitted for brevity as they are not relevant to the changes
     async def _start_websocket_manager(self):
         """Manages the WebSocket connection to the orchestrator."""
         while not self._shutdown_event.is_set():
-            for orchestrator in self._config.ORCHESTRATORS:
-                ws_url = orchestrator["url"].replace("http", "ws", 1) + "/_worker/ws"
+            # In multi-orchestrator mode, we currently only connect to the first one available
+            for _, client in self._clients:
                 try:
-                    if self._http_session:
-                        async with self._http_session.ws_connect(ws_url, headers=self._get_headers(orchestrator)) as ws:
-                            self._ws_connection = ws
-                            logger.info(f"WebSocket connection established to {ws_url}")
-                            await self._listen_for_commands()
-                except (ClientError, AsyncTimeoutError) as e:
-                    logger.warning(f"WebSocket connection to {ws_url} failed: {e}")
+                    ws = await client.connect_websocket()
+                    if ws:
+                        self._ws_connection = ws
+                        await self._listen_for_commands()
                 finally:
                     self._ws_connection = None
-                    logger.info(f"WebSocket connection to {ws_url} closed.")
                     await sleep(5)  # Reconnection delay
-            if not self._config.ORCHESTRATORS:
+            if not self._clients:
                 await sleep(5)
 
     async def _listen_for_commands(self):
@@ -513,7 +498,7 @@ class Worker:
                 if msg.type == WSMsgType.TEXT:
                     try:
                         command = msg.json()
-                        if command.get("type") == "cancel_task":
+                        if command.get("type") == COMMAND_CANCEL_TASK:
                             task_id = command.get("task_id")
                             if task_id in self._active_tasks:
                                 self._active_tasks[task_id].cancel()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: avtomatika-worker
-Version: 1.0b2
+Version: 1.0b3
 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
@@ -13,7 +13,7 @@ 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: obstore>=0.1
 Requires-Dist: aiofiles~=25.1.0
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
@@ -21,13 +21,14 @@ 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.
+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.
 
 ## Installation
 
@@ -472,7 +473,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.**
@@ -600,6 +601,7 @@ 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
 

avtomatika_worker-1.0b3.dist-info/RECORD

@@ -0,0 +1,14 @@
+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,,

avtomatika_worker-1.0b2.dist-info/RECORD

@@ -1,11 +0,0 @@
-avtomatika_worker/__init__.py,sha256=y_s5KlsgFu7guemZfjLVQ3Jzq7DyLG168-maVGwWRC4,334
-avtomatika_worker/config.py,sha256=v-2XGIcCIMr9S2SPAVKOMTpU8QSLeUm-udNOKWSxjQQ,5247
-avtomatika_worker/s3.py,sha256=ySwEOrP2ZslJ-Mg4_9vyxsnRzX0LIe78FmP2nlq8n9s,5930
-avtomatika_worker/task_files.py,sha256=ucjBuI78UmtMvfucTzDTNJ1g0KJaRIwyshRNTipIZSU,3351
-avtomatika_worker/types.py,sha256=MqXaX0NUatYDna3GgBWj73-WOT1EfaX1ei4i7eUsZR0,255
-avtomatika_worker/worker.py,sha256=sghM9Y8wwB4uJv_MPlQ3noBejDiwh4MYTZySNwjcZ3w,23968
-avtomatika_worker-1.0b2.dist-info/licenses/LICENSE,sha256=tqCjw9Y1vbU-hLcWi__7wQstLbt2T1XWPdbQYqCxuWY,1072
-avtomatika_worker-1.0b2.dist-info/METADATA,sha256=QdZKdrT-HP-u94gP7NVUt_TWGWJX6LIek2llHM1xj8Q,29184
-avtomatika_worker-1.0b2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-avtomatika_worker-1.0b2.dist-info/top_level.txt,sha256=d3b5BUeUrHM1Cn-cbStz-hpucikEBlPOvtcmQ_j3qAs,18
-avtomatika_worker-1.0b2.dist-info/RECORD,,