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

avtomatika_worker/config.py

@@ -12,56 +12,63 @@ 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.WORKER_PAYLOAD_DIR: str = getenv("WORKER_PAYLOAD_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]]:
         """
@@ -72,16 +79,20 @@ class WorkerConfig:
         if orchestrators_json:
             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:
         """Collects GPU information from environment variables.

avtomatika_worker/s3.py

@@ -0,0 +1,75 @@
+import asyncio
+import os
+from typing import Any
+from urllib.parse import urlparse
+
+import boto3
+from botocore.client import Config
+
+from .config import WorkerConfig
+
+
+class S3Manager:
+    """Handles S3 payload offloading."""
+
+    def __init__(self, config: WorkerConfig):
+        self._config = config
+        self._s3 = boto3.client(
+            "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 _process_s3_uri(self, uri: str) -> str:
+        """Downloads a file from S3 and returns the local path."""
+        parsed_url = urlparse(uri)
+        bucket_name = parsed_url.netloc
+        object_key = parsed_url.path.lstrip("/")
+        local_dir = self._config.WORKER_PAYLOAD_DIR
+        os.makedirs(local_dir, exist_ok=True)
+        local_path = os.path.join(local_dir, os.path.basename(object_key))
+
+        await asyncio.to_thread(self._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 to S3 and returns the S3 URI."""
+        bucket_name = self._config.S3_DEFAULT_BUCKET
+        object_key = os.path.basename(local_path)
+
+        await asyncio.to_thread(self._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]) -> 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)
+            if isinstance(item, dict):
+                return {k: await _process(v) for k, v in item.items()}
+            if isinstance(item, list):
+                return [await _process(i) for i in item]
+            return 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 os.path.exists(item) and item.startswith(self._config.WORKER_PAYLOAD_DIR):
+                return await self._upload_to_s3(item)
+            if isinstance(item, dict):
+                return {k: await _process(v) for k, v in item.items()}
+            if isinstance(item, list):
+                return [await _process(i) for i in item]
+            return item
+
+        return await _process(result)

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."""

avtomatika_worker/worker.py

@@ -1,5 +1,7 @@
 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
 from logging import getLogger
 from typing import Any, Callable
@@ -7,6 +9,15 @@ from typing import Any, Callable
 from aiohttp import ClientError, ClientSession, ClientTimeout, ClientWebSocketResponse, WSMsgType, web
 
 from .config import WorkerConfig
+from .s3 import S3Manager
+from .types import INVALID_INPUT_ERROR, PERMANENT_ERROR, TRANSIENT_ERROR, ParamValidationError
+
+try:
+    from pydantic import BaseModel, ValidationError
+
+    _PYDANTIC_INSTALLED = True
+except ImportError:
+    _PYDANTIC_INSTALLED = False
 
 # Logging setup
 logger = getLogger(__name__)
@@ -26,9 +37,11 @@ class Worker:
         task_type_limits: dict[str, int] | None = None,
         http_session: ClientSession | None = None,
         skill_dependencies: dict[str, list[str]] | None = None,
+        config: WorkerConfig | None = None,
     ):
-        self._config = WorkerConfig()
-        self._config.worker_type = worker_type  # Allow overriding worker_type
+        self._config = config or WorkerConfig()
+        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
 
@@ -44,12 +57,19 @@ class Worker:
         self._http_session = http_session
         self._session_is_managed_externally = http_session is not None
         self._ws_connection: ClientWebSocketResponse | None = None
-        self._headers = {"X-Worker-Token": self._config.worker_token}
+        # 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 ---
+        self._total_orchestrator_weight = 0
+        if self._config.ORCHESTRATORS:
+            for o in self._config.ORCHESTRATORS:
+                o["current_weight"] = 0
+                self._total_orchestrator_weight += o.get("weight", 1)
+
     def _validate_config(self):
         """Checks for unused task type limits and warns the user."""
         registered_task_types = {
@@ -98,7 +118,7 @@ class Worker:
         """
         Calculates the current worker state including status and available tasks.
         """
-        if self._current_load >= self._config.max_concurrent_tasks:
+        if self._current_load >= self._config.MAX_CONCURRENT_TASKS:
             return {"status": "busy", "supported_tasks": []}
 
         supported_tasks = []
@@ -118,9 +138,36 @@ 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:
+        """
+        Selects the next orchestrator using a smooth weighted round-robin algorithm.
+        """
+        if not self._config.ORCHESTRATORS:
+            return None
+
+        # The orchestrator with the highest current_weight is selected.
+        selected_orchestrator = None
+        highest_weight = -1
+
+        for o in self._config.ORCHESTRATORS:
+            o["current_weight"] += o["weight"]
+            if o["current_weight"] > highest_weight:
+                highest_weight = o["current_weight"]
+                selected_orchestrator = o
+
+        if selected_orchestrator:
+            selected_orchestrator["current_weight"] -= self._total_orchestrator_weight
+
+        return selected_orchestrator
+
     async def _debounced_heartbeat_sender(self):
         """Waits for the debounce delay then sends a heartbeat."""
-        await sleep(self._config.heartbeat_debounce_delay)
+        await sleep(self._config.HEARTBEAT_DEBOUNCE_DELAY)
         await self._send_heartbeats_to_all()
 
     def _schedule_heartbeat_debounce(self):
@@ -131,17 +178,18 @@ class Worker:
         # Schedule the new debounced call.
         self._debounce_task = create_task(self._debounced_heartbeat_sender())
 
-    async def _poll_for_tasks(self, orchestrator_url: str):
+    async def _poll_for_tasks(self, orchestrator: dict[str, Any]):
         """Polls a specific Orchestrator for new tasks."""
-        url = f"{orchestrator_url}/_worker/workers/{self._config.worker_id}/tasks/next"
+        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)
-            async with self._http_session.get(url, headers=self._headers, timeout=timeout) as resp:
+            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_url"] = orchestrator_url
+                    task_data["orchestrator"] = orchestrator
 
                     self._current_load += 1
                     task_handler_info = self._task_handlers.get(task_data["type"])
@@ -154,63 +202,125 @@ class Worker:
                     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)
+                    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)
+            await sleep(self._config.TASK_POLL_ERROR_DELAY)
 
     async def _start_polling(self):
-        print("Waiting for registration")
         """The main loop for polling tasks."""
         await self._registered_event.wait()
-        print("Polling started")
+
         while not self._shutdown_event.is_set():
             if self._get_current_state()["status"] == "busy":
-                await sleep(self._config.idle_poll_delay)
+                await sleep(self._config.IDLE_POLL_DELAY)
                 continue
 
-            if self._config.multi_orchestrator_mode == "ROUND_ROBIN":
-                orchestrator = self._config.orchestrators[self._round_robin_index]
-                await self._poll_for_tasks(orchestrator["url"])
-                self._round_robin_index = (self._round_robin_index + 1) % len(self._config.orchestrators)
+            if self._config.MULTI_ORCHESTRATOR_MODE == "ROUND_ROBIN":
+                orchestrator = self._get_next_orchestrator()
+                if orchestrator:
+                    await self._poll_for_tasks(orchestrator)
             else:
-                for orchestrator in self._config.orchestrators:
+                for orchestrator in self._config.ORCHESTRATORS:
                     if self._get_current_state()["status"] == "busy":
                         break
-                    await self._poll_for_tasks(orchestrator["url"])
+                    await self._poll_for_tasks(orchestrator)
 
             if self._current_load == 0:
-                await sleep(self._config.idle_poll_delay)
+                await sleep(self._config.IDLE_POLL_DELAY)
+
+    def _prepare_task_params(self, handler: Callable, params: dict[str, Any]) -> Any:
+        """
+        Inspects the handler's signature to validate and instantiate params.
+        Supports dict, dataclasses, and optional pydantic models.
+        """
+        sig = signature(handler)
+        params_annotation = sig.parameters.get("params").annotation
+
+        if params_annotation is sig.empty or params_annotation is dict:
+            return params
+
+        # Pydantic Model Validation
+        if _PYDANTIC_INSTALLED and isinstance(params_annotation, type) and issubclass(params_annotation, BaseModel):
+            try:
+                return params_annotation.model_validate(params)
+            except ValidationError as e:
+                raise ParamValidationError(str(e)) from e
+
+        # Dataclass Instantiation
+        if isinstance(params_annotation, type) and is_dataclass(params_annotation):
+            try:
+                # Filter unknown fields to prevent TypeError on dataclass instantiation
+                known_fields = {f.name for f in params_annotation.__dataclass_fields__.values()}
+                filtered_params = {k: v for k, v in params.items() if k in known_fields}
+
+                # Explicitly check for missing required fields
+                required_fields = [
+                    f.name
+                    for f in params_annotation.__dataclass_fields__.values()
+                    if f.default is Parameter.empty and f.default_factory is Parameter.empty
+                ]
+
+                missing_fields = [f for f in required_fields if f not in filtered_params]
+                if missing_fields:
+                    raise ParamValidationError(f"Missing required fields for dataclass: {', '.join(missing_fields)}")
+
+                return params_annotation(**filtered_params)
+            except (TypeError, ValueError) as e:
+                # TypeError for missing/extra args, ValueError from __post_init__
+                raise ParamValidationError(str(e)) from e
+
+        return params
 
     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_url = task_data.get("params", {}), task_data["orchestrator_url"]
+        params, orchestrator = task_data.get("params", {}), task_data["orchestrator"]
 
         result: dict[str, Any] = {}
         handler_data = self._task_handlers.get(task_name)
         task_type_for_limit = handler_data.get("type") if handler_data else None
 
+        result_sent = False  # Flag to track if result has been sent
+
         try:
-            if handler_data:
-                result = await handler_data["func"](
-                    params,
-                    task_id=task_id,
-                    job_id=job_id,
-                    priority=task_data.get("priority", 0),
-                    send_progress=self.send_progress,
-                    add_to_hot_cache=self.add_to_hot_cache,
-                    remove_from_hot_cache=self.remove_from_hot_cache,
-                )
-            else:
-                result = {"status": "failure", "error_message": f"Unsupported task: {task_name}"}
+            if not handler_data:
+                message = f"Unsupported task: {task_name}"
+                logger.warning(message)
+                result = {"status": "failure", "error": {"code": PERMANENT_ERROR, "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)
+                result_sent = True  # Mark result as sent
+                return
+
+            params = await self._s3_manager.process_params(params)
+            validated_params = self._prepare_task_params(handler_data["func"], params)
+
+            result = await handler_data["func"](
+                validated_params,
+                task_id=task_id,
+                job_id=job_id,
+                priority=task_data.get("priority", 0),
+                send_progress=self.send_progress,
+                add_to_hot_cache=self.add_to_hot_cache,
+                remove_from_hot_cache=self.remove_from_hot_cache,
+            )
+            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)}}
         except CancelledError:
+            logger.info(f"Task {task_id} was cancelled.")
             result = {"status": "cancelled"}
+            # We must re-raise the exception to be handled by the outer gather
+            raise
         except Exception as e:
-            result = {"status": "failure", "error": {"code": "TRANSIENT_ERROR", "message": str(e)}}
+            logger.exception(f"An unexpected error occurred while processing task {task_id}:")
+            result = {"status": "failure", "error": {"code": TRANSIENT_ERROR, "message": str(e)}}
         finally:
-            payload = {"job_id": job_id, "task_id": task_id, "worker_id": self._config.worker_id, "result": result}
-            await self._send_result(payload, orchestrator_url)
+            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)
             self._active_tasks.pop(task_id, None)
 
             self._current_load -= 1
@@ -218,14 +328,15 @@ 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_url: str):
+    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
-        for i in range(self._config.result_max_retries):
+        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=self._headers) as resp:
+                    async with self._http_session.post(url, json=payload, headers=headers) as resp:
                         if resp.status == 200:
                             return
             except ClientError as e:
@@ -233,43 +344,44 @@ class Worker:
             await sleep(delay * (2**i))
 
     async def _manage_orchestrator_communications(self):
-        print("Registering worker")
         """Registers the worker and sends heartbeats."""
         await self._register_with_all_orchestrators()
-        print("Worker registered")
+
         self._registered_event.set()
-        if self._config.enable_websockets:
+        if self._config.ENABLE_WEBSOCKETS:
             create_task(self._start_websocket_manager())
 
         while not self._shutdown_event.is_set():
             await self._send_heartbeats_to_all()
-            await sleep(self._config.heartbeat_interval)
+            await sleep(self._config.HEARTBEAT_INTERVAL)
 
     async def _register_with_all_orchestrators(self):
         """Registers the worker with all orchestrators."""
         state = self._get_current_state()
         payload = {
-            "worker_id": self._config.worker_id,
-            "worker_type": self._config.worker_type,
+            "worker_id": self._config.WORKER_ID,
+            "worker_type": self._config.WORKER_TYPE,
             "supported_tasks": state["supported_tasks"],
-            "max_concurrent_tasks": self._config.max_concurrent_tasks,
-            "installed_models": self._config.installed_models,
-            "hostname": self._config.hostname,
-            "ip_address": self._config.ip_address,
-            "resources": self._config.resources,
+            "max_concurrent_tasks": self._config.MAX_CONCURRENT_TASKS,
+            "cost_per_skill": self._config.COST_PER_SKILL,
+            "installed_models": self._config.INSTALLED_MODELS,
+            "hostname": self._config.HOSTNAME,
+            "ip_address": self._config.IP_ADDRESS,
+            "resources": self._config.RESOURCES,
         }
-        for orchestrator in self._config.orchestrators:
+        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._headers) as resp:
+                    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}")
 
     async def _send_heartbeats_to_all(self):
-        print("Sending heartbeats")
         """Sends heartbeat messages to all orchestrators."""
         state = self._get_current_state()
         payload = {
@@ -287,27 +399,27 @@ class Worker:
             if hot_skills:
                 payload["hot_skills"] = hot_skills
 
-        async def _send_single(orchestrator_url: str):
-            url = f"{orchestrator_url}/_worker/workers/{self._config.worker_id}"
+        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=self._headers) as resp:
+                    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}")
+                            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}")
+                logger.error(f"Error sending heartbeat to orchestrator {orchestrator['url']}: {e}")
 
-        await gather(*[_send_single(o["url"]) for o in self._config.orchestrators])
+        await gather(*[_send_single(o) for o in self._config.ORCHESTRATORS])
 
     async def main(self):
-        print("Main started")
         """The main asynchronous function."""
         self._validate_config()  # Validate config now that all tasks are registered
         if not self._http_session:
             self._http_session = ClientSession()
-        print("Starting comm task")
+
         comm_task = create_task(self._manage_orchestrator_communications())
-        print("Starting polling task")
+
         polling_task = create_task(self._start_polling())
         await self._shutdown_event.wait()
 
@@ -327,14 +439,13 @@ class Worker:
             run(self.main())
         except KeyboardInterrupt:
             self._shutdown_event.set()
-            run(sleep(1.5))
 
     async def _run_health_check_server(self):
         app = web.Application()
         app.router.add_get("/health", lambda r: web.Response(text="OK"))
         runner = web.AppRunner(app)
         await runner.setup()
-        site = web.TCPSite(runner, "0.0.0.0", self._config.worker_port)
+        site = web.TCPSite(runner, "0.0.0.0", self._config.WORKER_PORT)
         await site.start()
         await self._shutdown_event.wait()
         await runner.cleanup()
@@ -347,17 +458,16 @@ class Worker:
             run(_main_wrapper())
         except KeyboardInterrupt:
             self._shutdown_event.set()
-            run(sleep(1.5))
 
     # 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:
+            for orchestrator in self._config.ORCHESTRATORS:
                 ws_url = orchestrator["url"].replace("http", "ws", 1) + "/_worker/ws"
                 try:
                     if self._http_session:
-                        async with self._http_session.ws_connect(ws_url, headers=self._headers) as ws:
+                        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()
@@ -367,7 +477,7 @@ class Worker:
                     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._config.ORCHESTRATORS:
                 await sleep(5)
 
     async def _listen_for_commands(self):