manta-node 0.5b0.dev9__py3-none-any.whl

manta_node/node_orchestrator.py

@@ -0,0 +1,462 @@
+import asyncio
+import socket
+import struct
+import sys
+import uuid
+from logging import getLogger
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import blake3
+from aiodocker.exceptions import DockerError
+from docker.errors import DockerException
+from grpclib.exceptions import GRPCError
+from grpclib.server import Server
+
+from manta_common.build.common.informations import (
+    NodeOverview,
+    NodeStatus,
+    NodeStatusEnum,
+)
+from manta_common.build.common.system import Empty
+from manta_common.conversions import HEX_SIZE, ID, oid
+from manta_common.errors import MantaError
+from manta_common.traces import Tracer
+
+# NEW - use domain logic
+
+# In node_orchestrator.py
+from .infrastructure.config.node_config_manager import NodeConfigManager
+from .infrastructure.filesystem.dataset_manager import DatasetManager
+from .infrastructure.grpc.client import NodeClient
+from .infrastructure.grpc.local_servicer import LocalServicer
+from .infrastructure.grpc.world_servicer import WorldServicer
+from .infrastructure.metrics.metrics_collector import GenericMetricsCollector
+from .infrastructure.mqtt.command_handler import CommandHandler
+
+__all__ = ["Node"]
+
+
+def make_id(random_id: bool, alias: Optional[str] = None) -> ID:
+    """
+    Returns the node ID
+
+    Parameters
+    ----------
+    random_id : bool
+        If :code:`True`, the ID is randomized
+    alias : Optional[str]
+        Alias name of the node
+
+    Returns
+    -------
+    ID
+        Node ID
+    """
+    if alias is not None:
+        return ID(oid(alias))
+    elif random_id:
+        return ID(oid(uuid.uuid4().hex))
+    else:
+        return ID(
+            oid(blake3.blake3(struct.pack("Q", uuid.getnode())).hexdigest(HEX_SIZE))
+        )
+
+
+def find_available_port(start_port: int = 50051, max_attempts: int = 100) -> int:
+    """
+    Find an available port for the light service.
+
+    Parameters
+    ----------
+    start_port : int
+        The port to start searching from (default: 50051)
+    max_attempts : int
+        Maximum number of ports to try (default: 100)
+
+    Returns
+    -------
+    int
+        An available port number
+
+    Raises
+    ------
+    RuntimeError
+        If no available port is found after max_attempts
+    """
+    for port_offset in range(max_attempts):
+        port = start_port + port_offset
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            # Try to bind to the port
+            sock.bind(("", port))
+            sock.close()
+            return port
+        except OSError:
+            # Port is in use, try the next one
+            sock.close()
+            continue
+
+    raise RuntimeError(
+        f"Could not find an available port in range {start_port}-{start_port + max_attempts - 1}"
+    )
+
+
+class Node:
+    """
+    Node class runs on embedded device and executes tasks requested
+    by the Manager
+
+    It is composed of several components :
+
+    * the :code:`CommandHandler` which deploys tasks in containers \
+    and manages tasks given available resources
+    * the :code:`LocalServicer` to allow containers to access local data on the device
+    * the :code:`WorldServicer` to allow containers to access data stored in the Manager
+    * the :code:`MQTTClient` for gathering MQTT messages from the Manager
+    * the :code:`DatasetManager` for handling datasets and providing metadata
+
+    Parameters
+    ----------
+    config : Union[NodeConfigManager, str, Path]
+        NodeConfigManager instance or path to configuration file.
+        If string/Path is provided, a NodeConfigManager will be created.
+    """
+
+    __slots__ = [
+        "config_manager",
+        "config",
+        "node_id",
+        "node_client",
+        "task_runner",
+        "command_handler",
+        "server",
+        "tracer",
+        "cert_folder",
+        "dataset_manager",
+        "metrics_collector",
+        "command_loop",
+        "light_server_task",
+    ]
+
+    def __init__(self, config: Union[NodeConfigManager, str, Path]):
+        # Initialize configuration manager
+        if isinstance(config, NodeConfigManager):
+            self.config_manager = config
+        else:
+            # Create NodeConfigManager from path
+            self.config_manager = NodeConfigManager(config)
+
+        # Get configuration
+        self.config = self.config_manager.config
+
+        # Validate configuration
+        issues = self.config_manager.validate()
+        if issues:
+            raise MantaError(f"Configuration validation failed: {'; '.join(issues)}")
+
+        # Initialize tracer with logging configuration
+        self.tracer = Tracer(getLogger(__name__), None)
+
+        # Certificate folder setup (if TLS enabled)
+        if self.config.security.use_tls:
+            self.cert_folder = Path(self.config.security.cert_folder).expanduser()
+        else:
+            self.cert_folder = None
+
+        self.tracer.info(f"Configuration: {self.config.metadata.name}")
+
+        # Initialize the dataset manager with configuration
+        self.dataset_manager = DatasetManager(self.config.datasets.mappings)
+
+        # Validate dataset paths exist and are accessible
+        dataset_issues = self.dataset_manager.validate_datasets()
+        if dataset_issues:
+            error_msg = "Dataset validation failed:\n  - " + "\n  - ".join(
+                dataset_issues
+            )
+            self.tracer.error(error_msg)
+            raise MantaError(error_msg)
+
+        # Initialize the metrics collector
+        self.metrics_collector = GenericMetricsCollector()
+
+        # Initialize node client with configuration
+        self.node_client = NodeClient(
+            manager_host=self.config.network.manager_host,
+            manager_port=self.config.network.manager_port,
+            secured_token=self.config.identity.secured_token,
+            ssl_cert_folder=self.cert_folder,
+        )
+
+        # Initialize the command handler, the command loop and light server task
+        self.command_handler = None
+        self.command_loop = None
+        self.light_server_task = None
+
+        # Register configuration observer for runtime updates
+        self.config_manager.register_observer(self._on_config_change)
+
+    def _ensure_light_service_port(self) -> None:
+        """
+        Ensure light service port is assigned.
+
+        If not specified in configuration, automatically finds an available port.
+        Logs the assigned port for debugging purposes.
+        """
+        if self.config.network.light_service_port is None:
+            self.tracer.info(
+                "Light service port not specified, finding an available port..."
+            )
+            self.config.network.light_service_port = find_available_port()
+            self.tracer.info(
+                f"Automatically assigned light service port: {self.config.network.light_service_port}"
+            )
+
+    def _on_config_change(self, change: Dict[str, Any]) -> None:
+        """Handle configuration changes at runtime.
+
+        Parameters
+        ----------
+        change : Dict[str, Any]
+            Dictionary describing the configuration change
+        """
+        self.tracer.info(f"Configuration changed: {change}")
+
+        # Handle safe runtime updates
+        if change.get("action") == "update":
+            updates = change.get("updates", {})
+
+            # Update logging level if changed
+            if "logging.level" in updates:
+                import logging
+
+                logging.getLogger().setLevel(updates["logging.level"])
+                self.tracer.info(f"Updated logging level to {updates['logging.level']}")
+
+            # Update dataset mappings if changed
+            if "datasets.mappings" in updates:
+                self.dataset_manager.update_datasets(updates["datasets.mappings"])
+                self.tracer.info("Updated dataset mappings")
+
+            # Update resource limits if changed
+            if any(key.startswith("resources.") for key in updates):
+                if self.command_handler:
+                    # Command handler can update resource limits
+                    self.tracer.info("Resource limits updated")
+
+    async def check_manager_availability(self) -> None:
+        """
+        Check if the manager is available.
+
+        If secured mode is activated, this function will get the
+        certificates from the manager.
+        Then, it will check if the manager is available by calling
+        the `is_available` function of the node client.
+        If the manager is not available, it will log an error and
+        exit with a non-zero status code.
+
+        Raises
+        ------
+        SystemExit
+            If the manager is not available.
+        """
+        # Log endpoint details before attempting connection
+        security_mode = (
+            "TLS"
+            if (self.cert_folder is not None or self.config.network.manager_port == 443)
+            else "insecure"
+        )
+        endpoint = f"{security_mode}://{self.config.network.manager_host}:{self.config.network.manager_port}"
+
+        self.tracer.info(f"Checking manager availability at {endpoint}")
+
+        try:
+            # If secured mode, get the certificates
+            # if (
+            #     self.secured_token
+            # ):  # TODO: refactor with node bootstrap certificate and new identity service in node servicer
+            #     raise NotImplementedError("Secured mode is not implemented yet")
+            await self.node_client.is_available(Empty())
+            self.tracer.info(f"Manager is available at {endpoint}")
+        except (GRPCError, ConnectionRefusedError) as error:
+            message = f"Manager is not available at {endpoint}: {error}"
+            self.tracer.exception(message)
+
+            # Provide troubleshooting guidance
+            troubleshooting = (
+                f"\n[ERROR] {message}\n"
+                f"\nTroubleshooting steps:\n"
+                f"  1. Verify endpoint: {endpoint}\n"
+                f"  2. Check if manager services are running\n"
+                f"  3. Verify JWT_SECRET matches between node and manager\n"
+                f"  4. Check network connectivity to {self.config.network.manager_host}:{self.config.network.manager_port}\n"
+            )
+            sys.exit(troubleshooting)
+        except MantaError as manta_error:
+            message = f"Manager authentication failed at {endpoint}: {manta_error}"
+            self.tracer.exception(message)
+
+            # Provide specific guidance for authentication errors
+            troubleshooting = (
+                f"\n[ERROR] {message}\n"
+                f"\nAuthentication issue detected:\n"
+                f"  - Endpoint: {endpoint}\n"
+                f"  - This usually means JWT_SECRET mismatch between node and manager\n"
+                f"  - Verify secured_token in node config matches manager's JWT_SECRET\n"
+                f"  - Check token format: should be a valid JWT with 3 parts (header.payload.signature)\n"
+            )
+            sys.exit(troubleshooting)
+        except Exception as exc:
+            message = f"Unknown error connecting to {endpoint}: {exc}"
+            self.tracer.exception(message)
+            sys.exit(f"[ERROR] {message}")
+
+    async def start(self) -> None:
+        """
+        Start the Node:
+
+        - Register with the Manager
+        - Start the MQTT client
+        - Subscribe to topics
+        - Start the Light Server
+        - Start the MQTT client loop
+
+        Raises
+        ------
+        SystemExit
+            If MQTT client unreacheable or gRPC error occurs
+        """
+        await self.check_manager_availability()
+        try:
+            identification = NodeOverview(
+                self.metrics_collector.get_platform_info(),
+                self.dataset_manager.to_proto(),
+            )
+            registration = await self.node_client.register_node(identification)
+
+            # Initialize, start and prepare MQTT Client
+            self.tracer.debug(f"Registration: {registration}")
+            mqtt_broker_host = registration.mqtt_broker_host
+            mqtt_broker_port = registration.mqtt_broker_port
+            topics = registration.mqtt_topics
+            node_token = registration.node_token
+
+            # Update MQTT configuration with registration data
+            self.config.mqtt.broker_host = mqtt_broker_host
+            self.config.mqtt.broker_port = mqtt_broker_port
+            self.config.mqtt.topics = topics
+
+            # Save the node token to config file
+            if not node_token:
+                error_msg = "Manager did not return a node token during registration"
+                self.tracer.error(error_msg)
+                raise MantaError(error_msg)
+
+            self.tracer.info("Received node token from manager, updating config")
+            self.config.identity.secured_token = node_token
+            self.config_manager.save()
+
+            # Reinitialize the node client with the new token
+            self.node_client = NodeClient(
+                manager_host=self.config.network.manager_host,
+                manager_port=self.config.network.manager_port,
+                secured_token=node_token,
+                ssl_cert_folder=self.cert_folder,
+            )
+
+            # Ensure light service port is assigned
+            self._ensure_light_service_port()
+
+            # Initialize the command handler with configuration
+            self.tracer.info("Starting Command Handler")
+            self.command_handler = CommandHandler(
+                node_client=self.node_client,
+                mqtt_broker_host=mqtt_broker_host,
+                mqtt_broker_port=mqtt_broker_port,
+                light_service_port=self.config.network.light_service_port or 50052,
+                topics=topics,
+                dataset_manager=self.dataset_manager,
+                metrics_collector=self.metrics_collector,
+                stop_callback=self.stop,
+                ssl_cert_folder=self.cert_folder,
+                # Pass additional configuration
+                ping_interval=self.config.mqtt.ping_interval,
+                max_concurrent_tasks=self.config.tasks.max_concurrent,
+                task_timeout=self.config.tasks.task_timeout,
+            )
+
+            # Start the server for docker communications through
+            # manta light with manta node
+            self.tracer.info(
+                f"Starting Light Server at {self.config.network.light_service_host}:"
+                f"{self.config.network.light_service_port}"
+            )
+            local_servicer = LocalServicer(
+                self.node_client,
+                self.dataset_manager,
+            )
+            world_servicer = WorldServicer(
+                self.node_client,
+                self.command_handler.tasks,
+            )
+            self.server = Server([local_servicer, world_servicer])
+            await self.server.start(
+                self.config.network.light_service_host,
+                self.config.network.light_service_port,
+            )
+
+            # Create asyncio tasks and start them
+            self.command_loop = asyncio.create_task(self.command_handler.main())
+            self.light_server_task = asyncio.create_task(self.server.wait_closed())
+            await asyncio.gather(self.command_loop, self.light_server_task)
+        except ConnectionRefusedError:
+            self.tracer.exception("Unable to start and connect to MQTT Message Broker")
+            sys.exit("[ERROR] Unable to start and connect to MQTT Message Broker")
+        except GRPCError:
+            self.tracer.exception("Unable to register to the Manager")
+            sys.exit("[ERROR] Unable to register to the Manager")
+        except OSError as os_error:
+            self.tracer.exception(f"OS error: {os_error}")
+            await self.stop()
+            sys.exit(f"[ERROR] OS error: {os_error}")
+        except (DockerException, DockerError) as docker_error:
+            self.tracer.exception(f"Docker is not available: {docker_error}")
+            await self.stop()
+            sys.exit(f"[ERROR] Docker is not available: {docker_error}")
+        except Exception as exc:
+            self.tracer.exception(f"Unknown error: {exc}")
+            await self.stop()
+            sys.exit(f"[ERROR] Unknown error: {exc}")
+
+    async def stop(self) -> None:
+        """
+        Stop the node
+        """
+        self.tracer.info("Stopping node")
+        if (
+            self.command_handler is not None
+            and self.command_loop is not None
+            and not self.command_loop.cancelled()
+        ):
+            self.tracer.info("Publishing node status to manager")
+            await self.command_handler.client.publish(
+                "manager/set_node_status",
+                NodeStatus(status=NodeStatusEnum.DISCONNECTED).SerializeToString(),
+            )
+
+        if self.command_handler is not None:
+            self.tracer.info("Cleaning up command handler")
+            await self.command_handler.clean()
+
+        if self.command_loop is not None:
+            self.tracer.info("Stopping command loop")
+            self.command_loop.cancel()
+
+        if self.light_server_task is not None:
+            self.tracer.info("Stopping light server task")
+            self.light_server_task.cancel()
+
+        # if self.server is not None:
+        #     self.tracer.info("Stopping light server")
+        #     self.server.close()