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

avtomatika_worker/__init__.py

@@ -2,9 +2,10 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
+from .task_files import TaskFiles
 from .worker import Worker
 
-__all__ = ["Worker"]
+__all__ = ["Worker", "TaskFiles"]
 
 try:
     __version__ = version("avtomatika-worker")

avtomatika_worker/config.py

@@ -12,94 +12,104 @@ class WorkerConfig:
 
     def __init__(self):
         # --- Basic worker information ---
-        self.worker_id: str = getenv("WORKER_ID", f"worker-{uuid4()}")
-        self.worker_type: str = getenv("WORKER_TYPE", "generic-cpu-worker")
-        self.worker_port: int = int(getenv("WORKER_PORT", "8083"))
-        self.hostname: str = gethostname()
+        self.WORKER_ID: str = getenv("WORKER_ID", f"worker-{uuid4()}")
+        self.WORKER_TYPE: str = getenv("WORKER_TYPE", "generic-cpu-worker")
+        self.WORKER_PORT: int = int(getenv("WORKER_PORT", "8083"))
+        self.HOSTNAME: str = gethostname()
         try:
-            self.ip_address: str = gethostbyname(self.hostname)
+            self.IP_ADDRESS: str = gethostbyname(self.HOSTNAME)
         except gaierror:
-            self.ip_address: str = "127.0.0.1"
+            self.IP_ADDRESS: str = "127.0.0.1"
 
         # --- Orchestrator settings ---
-        self.orchestrators: list[dict[str, Any]] = self._get_orchestrators_config()
+        self.ORCHESTRATORS: list[dict[str, Any]] = self._get_orchestrators_config()
 
         # --- Security ---
-        self.worker_token: str = getenv(
+        self.WORKER_TOKEN: str = getenv(
             "WORKER_INDIVIDUAL_TOKEN",
             getenv("WORKER_TOKEN", "your-secret-worker-token"),
         )
 
         # --- Resources and performance ---
-        self.cost_per_second: float = float(getenv("WORKER_COST_PER_SECOND", "0.01"))
-        self.max_concurrent_tasks: int = int(getenv("MAX_CONCURRENT_TASKS", "10"))
-        self.resources: dict[str, Any] = {
+        self.COST_PER_SKILL: dict[str, float] = self._load_json_from_env("COST_PER_SKILL", default={})
+        self.MAX_CONCURRENT_TASKS: int = int(getenv("MAX_CONCURRENT_TASKS", "10"))
+        self.RESOURCES: dict[str, Any] = {
             "cpu_cores": int(getenv("CPU_CORES", "4")),
             "gpu_info": self._get_gpu_info(),
         }
 
         # --- Installed software and models (read as JSON strings) ---
-        self.installed_software: dict[str, str] = self._load_json_from_env(
+        self.INSTALLED_SOFTWARE: dict[str, str] = self._load_json_from_env(
             "INSTALLED_SOFTWARE",
             default={"python": "3.9"},
         )
-        self.installed_models: list[dict[str, str]] = self._load_json_from_env(
+        self.INSTALLED_MODELS: list[dict[str, str]] = self._load_json_from_env(
             "INSTALLED_MODELS",
             default=[],
         )
 
+        # --- S3 Settings for payload offloading ---
+        self.TASK_FILES_DIR: str = getenv("TASK_FILES_DIR", "/tmp/payloads")
+        self.S3_ENDPOINT_URL: str | None = getenv("S3_ENDPOINT_URL")
+        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")
+
         # --- Tuning parameters ---
-        self.heartbeat_interval: float = float(getenv("HEARTBEAT_INTERVAL", "15"))
-        self.result_max_retries: int = int(getenv("RESULT_MAX_RETRIES", "5"))
-        self.result_retry_initial_delay: float = float(
+        self.HEARTBEAT_INTERVAL: float = float(getenv("HEARTBEAT_INTERVAL", "15"))
+        self.RESULT_MAX_RETRIES: int = int(getenv("RESULT_MAX_RETRIES", "5"))
+        self.RESULT_RETRY_INITIAL_DELAY: float = float(
             getenv("RESULT_RETRY_INITIAL_DELAY", "1.0"),
         )
-        self.heartbeat_debounce_delay: float = float(getenv("WORKER_HEARTBEAT_DEBOUNCE_DELAY", 0.1))
-        self.task_poll_timeout: float = float(getenv("TASK_POLL_TIMEOUT", "30"))
-        self.task_poll_error_delay: float = float(
+        self.HEARTBEAT_DEBOUNCE_DELAY: float = float(getenv("WORKER_HEARTBEAT_DEBOUNCE_DELAY", 0.1))
+        self.TASK_POLL_TIMEOUT: float = float(getenv("TASK_POLL_TIMEOUT", "30"))
+        self.TASK_POLL_ERROR_DELAY: float = float(
             getenv("TASK_POLL_ERROR_DELAY", "5.0"),
         )
-        self.idle_poll_delay: float = float(getenv("IDLE_POLL_DELAY", "0.01"))
-        self.enable_websockets: bool = getenv("WORKER_ENABLE_WEBSOCKETS", "false").lower() == "true"
-        self.multi_orchestrator_mode: str = getenv("MULTI_ORCHESTRATOR_MODE", "FAILOVER")
+        self.IDLE_POLL_DELAY: float = float(getenv("IDLE_POLL_DELAY", "0.01"))
+        self.ENABLE_WEBSOCKETS: bool = getenv("WORKER_ENABLE_WEBSOCKETS", "false").lower() == "true"
+        self.MULTI_ORCHESTRATOR_MODE: str = getenv("MULTI_ORCHESTRATOR_MODE", "FAILOVER")
 
     def _get_orchestrators_config(self) -> list[dict[str, Any]]:
         """
         Loads orchestrator configuration from the ORCHESTRATORS_CONFIG environment variable.
         For backward compatibility, if it is not set, it uses ORCHESTRATOR_URL.
         """
-        orchestrators_json = getenv("ORCHESTRATORS_CONFIG")
-        if orchestrators_json:
+        if orchestrators_json := getenv("ORCHESTRATORS_CONFIG"):
             try:
                 orchestrators = loads(orchestrators_json)
+                if getenv("ORCHESTRATOR_URL"):
+                    print("Info: Both ORCHESTRATORS_CONFIG and ORCHESTRATOR_URL are set. Using ORCHESTRATORS_CONFIG.")
                 for o in orchestrators:
                     if "priority" not in o:
                         o["priority"] = 10
+                    if "weight" not in o:
+                        o["weight"] = 1
                 orchestrators.sort(key=lambda x: (x.get("priority", 10), x.get("url")))
                 return orchestrators
             except JSONDecodeError:
                 print("Warning: Could not decode JSON from ORCHESTRATORS_CONFIG. Falling back to default.")
 
         orchestrator_url = getenv("ORCHESTRATOR_URL", "http://localhost:8080")
-        return [{"url": orchestrator_url, "priority": 1}]
+        return [{"url": orchestrator_url, "priority": 1, "weight": 1}]
 
-    def _get_gpu_info(self) -> dict[str, Any] | None:
+    @staticmethod
+    def _get_gpu_info() -> dict[str, Any] | None:
         """Collects GPU information from environment variables.
         Returns None if GPU is not configured.
         """
-        gpu_model = getenv("GPU_MODEL")
-        if not gpu_model:
+        if gpu_model := getenv("GPU_MODEL"):
+            return {
+                "model": gpu_model,
+                "vram_gb": int(getenv("GPU_VRAM_GB", "0")),
+            }
+        else:
             return None
 
-        return {
-            "model": gpu_model,
-            "vram_gb": int(getenv("GPU_VRAM_GB", "0")),
-        }
-
-    def _load_json_from_env(self, key: str, default: Any) -> Any:
+    @staticmethod
+    def _load_json_from_env(key: str, default: Any) -> Any:
         """Safely loads a JSON string from an environment variable."""
-        value = getenv(key)
-        if value:
+        if value := getenv(key):
             try:
                 return loads(value)
             except JSONDecodeError:

avtomatika_worker/s3.py

@@ -0,0 +1,141 @@
+from asyncio import gather, to_thread
+from os import walk
+from os.path import basename, dirname, join, relpath
+from shutil import rmtree
+from typing import Any
+from urllib.parse import urlparse
+
+from aioboto3 import Session
+from aiofiles.os import makedirs
+from aiofiles.ospath import exists, isdir
+from botocore.client import Config
+
+from .config import WorkerConfig
+
+
+class S3Manager:
+    """Handles S3 payload offloading."""
+
+    def __init__(self, config: WorkerConfig):
+        self._config = config
+        self._session = Session()
+
+    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,
+            "aws_access_key_id": self._config.S3_ACCESS_KEY,
+            "aws_secret_access_key": self._config.S3_SECRET_KEY,
+            "config": Config(signature_version="s3v4"),
+        }
+
+    async def cleanup(self, task_id: str):
+        """Removes the task-specific payload directory."""
+        task_dir = join(self._config.TASK_FILES_DIR, task_id)
+        if await exists(task_dir):
+            await to_thread(lambda: rmtree(task_dir, ignore_errors=True))
+
+    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("/")
+
+        # 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)
+
+        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
+
+                        # 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))
+
+                if tasks:
+                    await gather(*tasks)
+                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)
+            return local_path
+
+    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
+
+        async with self._session.client(**self._get_client_args()) as s3:
+            # 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():
+                    files_to_upload = []
+                    for root, _, files in walk(local_path):
+                        for file in files:
+                            f_path = join(root, file)
+                            rel = relpath(f_path, local_path)
+                            files_to_upload.append((f_path, f"{s3_prefix}{rel}"))
+                    return files_to_upload
+
+                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))
+
+                if tasks:
+                    await gather(*tasks)
+
+                return f"s3://{bucket_name}/{s3_prefix}"
+
+            # 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}"
+
+    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."""
+        if not self._config.S3_ENDPOINT_URL:
+            return params
+
+        async def _process(item: Any) -> Any:
+            if isinstance(item, str) and item.startswith("s3://"):
+                return await self._process_s3_uri(item, task_id)
+            if isinstance(item, dict):
+                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)
+
+    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."""
+        if not self._config.S3_ENDPOINT_URL:
+            return result
+
+        async def _process(item: Any) -> Any:
+            if isinstance(item, str) and item.startswith(self._config.TASK_FILES_DIR):
+                return await self._upload_to_s3(item) if await exists(item) else item
+            if isinstance(item, dict):
+                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)

avtomatika_worker/task_files.py

@@ -0,0 +1,97 @@
+from contextlib import asynccontextmanager
+from os.path import dirname, join
+from typing import AsyncGenerator
+
+from aiofiles import open as aiopen
+from aiofiles.os import listdir, makedirs
+from aiofiles.ospath import exists as aio_exists
+
+
+class TaskFiles:
+    """
+    A helper class for managing task-specific files.
+    Provides asynchronous lazy directory creation and high-level file operations
+    within an isolated workspace for each task.
+    """
+
+    def __init__(self, task_dir: str):
+        """
+        Initializes TaskFiles with a specific task directory.
+        The directory is not created until needed.
+        """
+        self._task_dir = task_dir
+
+    async def get_root(self) -> str:
+        """
+        Asynchronously returns the root directory for the task.
+        Creates the directory on disk if it doesn't exist.
+        """
+        await makedirs(self._task_dir, exist_ok=True)
+        return self._task_dir
+
+    async def path_to(self, filename: str) -> str:
+        """
+        Asynchronously returns an absolute path for a file within the task directory.
+        Guarantees that the task root directory exists.
+        """
+        root = await self.get_root()
+        return join(root, filename)
+
+    @asynccontextmanager
+    async def open(self, filename: str, mode: str = "r") -> AsyncGenerator:
+        """
+        An asynchronous context manager to open a file within the task directory.
+        Automatically creates the task root and any necessary subdirectories.
+
+        Args:
+            filename: Name or relative path of the file.
+            mode: File opening mode (e.g., 'r', 'w', 'a', 'rb', 'wb').
+        """
+        path = await self.path_to(filename)
+        # Ensure directory for the file itself exists if filename contains subdirectories
+        file_dir = dirname(path)
+        if file_dir != self._task_dir:
+            await makedirs(file_dir, exist_ok=True)
+
+        async with aiopen(path, mode) as f:
+            yield f
+
+    async def read(self, filename: str, mode: str = "r") -> str | bytes:
+        """
+        Asynchronously reads the entire content of a file.
+
+        Args:
+            filename: Name of the file to read.
+            mode: Mode to open the file in (defaults to 'r').
+        """
+        async with self.open(filename, mode) as f:
+            return await f.read()
+
+    async def write(self, filename: str, data: str | bytes, mode: str = "w") -> None:
+        """
+        Asynchronously writes data to a file. Creates or overwrites the file by default.
+
+        Args:
+            filename: Name of the file to write.
+            data: Content to write (string or bytes).
+            mode: Mode to open the file in (defaults to 'w').
+        """
+        async with self.open(filename, mode) as f:
+            await f.write(data)
+
+    async def list(self) -> list[str]:
+        """
+        Asynchronously lists all file and directory names within the task root.
+        """
+        root = await self.get_root()
+        return await listdir(root)
+
+    async def exists(self, filename: str) -> bool:
+        """
+        Asynchronously checks if a specific file or directory exists in the task root.
+        """
+        path = join(self._task_dir, filename)
+        return await aio_exists(path)
+
+    def __repr__(self):
+        return f"<TaskFiles root='{self._task_dir}'>"

avtomatika_worker/types.py

@@ -2,3 +2,7 @@
 TRANSIENT_ERROR = "TRANSIENT_ERROR"
 PERMANENT_ERROR = "PERMANENT_ERROR"
 INVALID_INPUT_ERROR = "INVALID_INPUT_ERROR"
+
+
+class ParamValidationError(Exception):
+    """Custom exception for parameter validation errors."""