manta-node 0.5b0__py3-none-any.whl

manta_node/infrastructure/mqtt/command_handler.py

@@ -0,0 +1,450 @@
+import asyncio
+import io
+from datetime import datetime, timezone
+from logging import getLogger
+from pathlib import Path
+from typing import AsyncIterable, Awaitable, Callable, List, Optional, cast
+
+import aiomqtt
+
+from manta_common.base_mqtt import MqttBase
+from manta_common.build.common.mqtt import MqMetricsSnapshot, MqTaskUpdate
+from manta_common.build.common.system import Logs
+from manta_common.build.common.tasks import MqttTask, TaskStatus, TaskUpdate
+from manta_common.const import CHUNK_SIZE
+from manta_common.conversions import ID
+from manta_common.errors import (
+    MantaError,
+    MantaGRPCError,
+    MantaMQTTError,
+    MantaNodeError,
+)
+from manta_common.traces import Tracer
+from ...task_manager import TaskManager, LOG_FLUSH_INTERVAL
+from ..filesystem.dataset_manager import DatasetManager
+from ..grpc.client import NodeClient
+from ..metrics.collector import Collector
+from ..metrics.metrics_collector import MetricsCollectorBase
+
+TIMEOUT = 1  # s
+
+
+class CommandHandler(TaskManager, Collector, MqttBase):
+    """
+    CommandHandler class to handle commands from the Manager
+
+    Parameters
+    ----------
+    node_id: ID
+        Node ID
+    mqtt_client: MQTTClient
+        MQTT client
+    node_client: NodeClient
+        Node client
+    mqtt_broker_host : str
+        MQTT broker host
+    mqtt_broker_port : int
+        MQTT broker port
+    topics : List[str]
+        List of topics
+    light_service_port: int
+        Light service port
+    dataset_manager: DatasetManager
+        Dataset manager
+    metrics_collector: MetricsCollectorBase
+        Metrics collector
+    stop_callback: Callable[[], Awaitable[None]]
+        Stop callback
+    ssl_cert_folder : Optional[Path]
+        SSL Certification folder
+    """
+
+    def __init__(
+        self,
+        node_client: NodeClient,
+        mqtt_broker_host: str,
+        mqtt_broker_port: int,
+        light_service_port: int,
+        topics: List[str],
+        dataset_manager: DatasetManager,
+        metrics_collector: MetricsCollectorBase,
+        stop_callback: Callable[[], Awaitable[None]],
+        ssl_cert_folder: Optional[Path] = None,
+        ping_interval: int = TIMEOUT,
+        max_concurrent_tasks: int = 2,
+        task_timeout: int = 3600,
+    ) -> None:
+        self.tracer = Tracer(getLogger(__name__))
+        self.node_client = node_client
+        self.dataset_manager = dataset_manager
+        self.metrics_collector = metrics_collector
+        self.stop_callback = stop_callback
+        self.ping_interval = ping_interval
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.task_timeout = task_timeout
+
+        MqttBase.__init__(
+            self,
+            mqtt_broker_host=mqtt_broker_host,
+            mqtt_broker_port=mqtt_broker_port,
+            topics=topics,
+            client_id=ID(node_client.metadata.get("authorization", "node_client")),
+            ssl_cert_folder=ssl_cert_folder,
+        )
+
+        self.tracer.debug("Trying to connect to docker daemon...")
+        TaskManager.__init__(self, light_service_port)
+        self.tracer.info("Connected to docker daemon !")
+
+    async def clean(self):  # TODO : use __del__
+        """Delete the temporary folder and stop all containers"""
+        self.tracer.debug("Cleaning up the temporary folder")
+        self.temp_folder.cleanup()
+
+        self.tracer.info("Stopping all tasks")
+        for task_progression in self.tasks:
+            try:
+                task_id = task_progression.task_id
+                swarm_id = task_progression.swarm_id
+
+                task_progression.unlink()
+                await self.stop_container(swarm_id, task_id)
+                await self.remove_container(swarm_id, task_id)
+            except Exception:
+                self.tracer.exception(
+                    f"Error when stopping task {task_progression.task_id}"
+                )
+
+        self.tracer.info("Closing Docker client")
+        try:
+            await self.async_client.close()
+        except Exception:
+            self.tracer.exception("Error when closing Docker client")
+
+    async def main(self):
+        """
+        Override MqttBase.main to run task checker as a background task
+        alongside the MQTT message loop.
+        """
+        self.tracer.debug("Starting background task checker")
+        checker = asyncio.create_task(self._ping_loop())
+        try:
+            await MqttBase.main(self)
+        finally:
+            checker.cancel()
+            try:
+                await checker
+            except asyncio.CancelledError:
+                pass
+
+    async def loop_tasks(self):
+        """
+        MQTT message loop (subscribe + publish).
+        Task checking runs independently in the background via main().
+
+        Overrides base class to fix a race condition: the background
+        _ping_loop can queue messages during subscribe processing,
+        causing publish_task (queue.get()) to consume and remove them
+        from the queue. If the base class processes subscribe first
+        without also handling the completed publish_task, the message
+        is silently dropped. See issue #285.
+        """
+        while True:
+            try:
+                subscribe_task = asyncio.create_task(self.client.messages.__anext__())
+                publish_task = asyncio.create_task(self.queue.get())
+                done, _ = await asyncio.wait(
+                    (subscribe_task, publish_task),
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+            except asyncio.CancelledError:
+                subscribe_task.cancel()
+                publish_task.cancel()
+                await asyncio.gather(
+                    subscribe_task, publish_task, return_exceptions=True
+                )
+                raise
+
+            try:
+                if subscribe_task in done and publish_task in done:
+                    try:
+                        await self._subscribe_from_task(subscribe_task)
+                    finally:
+                        await self._publish_from_task(publish_task)
+                elif subscribe_task in done:
+                    await self._subscribe_from_task(subscribe_task)
+                    if publish_task.done():
+                        await self._publish_from_task(publish_task)
+                elif publish_task in done:
+                    await self._publish_from_task(publish_task)
+                else:
+                    self.tracer.warning("Disconnected during message iteration")
+                    raise aiomqtt.MqttError("Disconnected during message iteration")
+            finally:
+                if not subscribe_task.done():
+                    subscribe_task.cancel()
+                if not publish_task.done():
+                    publish_task.cancel()
+
+    async def _sweep_orphans(self) -> None:
+        """Kill Docker containers with manta labels that are no longer tracked.
+
+        When a ``stop_swarm`` MQTT message is lost (connection drop) the
+        Docker container stays alive and blocks port 8080 on subsequent
+        deploys.  This runs periodically in the background to clean up.
+        """
+        try:
+            docker_containers = await self.async_client.containers.list(
+                all=False,
+                filters={"label": ["com.docker.compose.project=manta_tasks"]},
+            )
+        except Exception:
+            return
+        tracked = set(self.tasks._env_ids) | self._launching
+
+        for dc in docker_containers:
+            labels = dc.get("Labels", {}) or {}
+            swarm_xid = labels.get("manta.swarm-id")
+            task_xid = labels.get("manta.task-id")
+            if swarm_xid is None or task_xid is None:
+                continue
+            tid_pair = (ID(swarm_xid).tid, ID(task_xid).tid)
+            if tid_pair in tracked:
+                continue
+
+            cid = dc["Id"][:12]
+            self.tracer.info(
+                f"Orphan sweep: killing container {cid} "
+                f"({swarm_xid[:8]}/{task_xid[:8]})"
+            )
+            try:
+                container = await self.async_client.containers.get(dc["Id"])
+                await container.kill()
+                await container.delete()
+            except Exception:
+                self.tracer.exception(f"Failed to clean up orphan container {cid}")
+
+    async def _ping_loop(self):
+        """
+        Background loop that periodically checks task statuses and
+        sends heartbeats. Runs independently of the MQTT connection state.
+        """
+        iteration = 0
+        while True:
+            iteration += 1
+            if iteration % 10 == 0:
+                self.tracer.info(
+                    f"Ping loop alive: iteration {iteration} "
+                    f"(tasks={len(self.tasks._env_ids)}, queue={len(self.task_queue)})"
+                )
+            if iteration % 15 == 0:
+                try:
+                    await self._sweep_orphans()
+                except Exception:
+                    self.tracer.exception("Orphan sweep failed")
+            try:
+                await self._ping()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                self.tracer.exception("Unhandled error in ping loop, continuing...")
+
+    async def _ping(self):
+        """
+        Ping the Manager
+        """
+        await asyncio.sleep(self.ping_interval)
+        updated = await self.check_tasks()
+        if len(updated) > 0:
+            for task_progression in updated:
+                if task_progression.task_status == TaskStatus.STOPPED:
+                    self.tracer.info(
+                        f"Task {task_progression.task_id} STOPPED, collecting logs"
+                    )
+                    try:
+                        date, processed_logs = await self.get_task_logs(
+                            task_progression
+                        )
+                        self.tracer.info(
+                            f"Sending logs to the Manager for task {task_progression.task_id}"
+                        )
+                        await self.send_task_logs(date, processed_logs)
+                    except MantaGRPCError as e:
+                        self.tracer.exception(
+                            f"gRPCError when updating finished tasks: {e}"
+                        )
+                        task_progression.task_status = TaskStatus.FAILED
+                    except MantaError as manta_error:
+                        self.tracer.exception("Error updating finished tasks.")
+                        manta_error.metadata = self.tracer.metadata
+                        task_progression.task_status = TaskStatus.FAILED
+                        await self.node_client.send_error(manta_error)
+                    except Exception as exc:
+                        message = f"Error updating finished tasks: {repr(exc)}"
+                        self.tracer.exception(message)
+                        task_progression.task_status = TaskStatus.FAILED
+                        await self.node_client.send_error(
+                            MantaNodeError(message, metadata=self.tracer.metadata)
+                        )
+                else:
+                    self.tracer.info(
+                        f"Publishing status for task {task_progression.task_id}: "
+                        f"{task_progression.task_status}"
+                    )
+                    try:
+                        await self.publish_task_update(task_progression.task_update)
+                        self.tracer.info(
+                            f"Status published for task {task_progression.task_id}: "
+                            f"{task_progression.task_status}"
+                        )
+                    except Exception as exc:
+                        message = f"Error updating active tasks: {repr(exc)}"
+                        self.tracer.exception("Error updating active tasks.")
+                        await self.node_client.send_error(
+                            MantaNodeError(message, metadata=self.tracer.metadata)
+                        )
+        else:
+            await self.ping()
+
+        # Option B: periodic log flush for active RUNNING tasks
+        now = datetime.now(timezone.utc)
+        for tp in self.tasks.active_tasks():
+            if tp.task_status != TaskStatus.RUNNING:
+                continue
+            elapsed = (now - tp.last_log_collected).total_seconds()
+            if elapsed >= LOG_FLUSH_INTERVAL:
+                try:
+                    await self._flush_task_logs(tp)
+                except Exception:
+                    self.tracer.exception(
+                        f"Periodic log flush failed for task {tp.task_id}"
+                    )
+
+    async def ping(self):
+        """
+        Ping the Manager
+        """
+        try:
+            metrics = await self.metrics_collector.collect_all_metrics()
+            await self.publish_message(
+                "manager/set_system_summary",
+                MqMetricsSnapshot(
+                    token=self.node_client.secured_token, snapshot=metrics
+                ),
+            )
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error when pinging the Manager: {e}")
+        except MantaError as e:
+            self.tracer.exception(f"Error when pinging the Manager: {e}")
+            await self.node_client.send_error(e)
+
+    async def process_message(self, message: aiomqtt.Message):
+        """
+        Handle messages from MQTT
+
+        Parameters
+        ----------
+        message : aiomqtt.Message
+            Message
+        """
+        try:
+            self.tracer.clean_metadata()
+            key = str(message.topic).split("/")[-1]
+            payload = cast(bytes, message.payload)
+
+            if key == "deploy_task":
+                await self.deploy_task(payload)
+            elif key == "stop_task":
+                task_update = await self.stop_task(payload)
+                if task_update is not None:
+                    await self.publish_task_update(task_update)
+            elif key == "stop_swarm":
+                await self.stop_swarm(payload)
+            elif key == "collect_information":
+                await self.collect_information(payload)
+            elif key == "stop_node":
+                await self.stop_node(payload)
+            else:
+                self.tracer.error(f"Not allowed method: {key}")
+                raise MantaMQTTError(
+                    f"Not allowed method: {key}", metadata=self.tracer.metadata
+                )
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error sending message to the Manager: {e}")
+        except MantaError as manta_error:
+            self.tracer.exception(f"Error when handling message: {manta_error}")
+            await self.node_client.send_error(manta_error)
+
+    async def publish_task_update(self, task_update: TaskUpdate):
+        """
+        Publish a task update to the Manager through MQTT
+        Message broker
+
+        Parameters
+        ----------
+        task_update : TaskUpdate
+            Task Update
+        """
+        await self.publish_message(
+            "manager/update_task_status",
+            MqTaskUpdate(token=self.node_client.secured_token, update=task_update),
+        )
+
+    async def send_task_logs(self, date: datetime, processed_logs: bytes):
+        """
+        Send task logs to the Manager
+
+        Parameters
+        ----------
+        date : datetime
+            Date when logs where created
+        processed_logs : bytes
+            Logs prepared from tracer
+        """
+
+        async def chunk_logs(content: bytes) -> AsyncIterable[Logs]:
+            data_stream = io.BytesIO(content)
+            while chunk := data_stream.read(CHUNK_SIZE):
+                yield Logs(content=chunk, timestamp=date)
+
+        await self.node_client.add_task_logs(chunk_logs(processed_logs))
+
+    async def stop_node(self, _payload: bytes):
+        """
+        Stop the node and prepare for graceful shutdown
+
+        Parameters
+        ----------
+        payload : bytes
+            Node ID to stop
+        """
+        await self.stop_callback()
+
+    async def on_task_setup_failed(self, task: MqttTask, exc: Exception):
+        """
+        Publish a FAILED task update when task setup fails.
+
+        Parameters
+        ----------
+        task : MqttTask
+            Task that failed to setup
+        exc : Exception
+            The exception that caused the failure
+        """
+        task_update = TaskUpdate(
+            swarm_id=task.swarm_id,
+            node_id=self.tracer.metadata.node_id,
+            task_id=task.task_id,
+            iteration=task.iteration,
+            circular=task.circular,
+            task_status=TaskStatus.FAILED,
+        )
+        try:
+            await self.publish_task_update(task_update)
+            self.tracer.info(
+                f"Published FAILED status for task {task.task_id} after setup failure"
+            )
+        except Exception as publish_exc:
+            self.tracer.exception(
+                f"Failed to publish FAILED status for task {task.task_id}: {repr(publish_exc)}"
+            )