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

manta_node/task_manager.py

@@ -0,0 +1,519 @@
+import asyncio
+from collections import deque
+from datetime import datetime
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import List, Optional, Tuple
+
+from manta_common.build.common.system import EnvId, PairId
+from manta_common.build.common.tasks import MqttTask, TaskStatus, TaskUpdate
+from manta_common.conversions import ID
+from manta_common.errors import (
+    MantaError,
+    MantaGRPCError,
+    MantaMQTTError,
+    MantaNodeError,
+)
+from manta_common.traces import Tracer
+from .domain.task_lifecycle import TaskLifecycleService
+from .infrastructure.container.manager import ContainerManager
+from .infrastructure.grpc.client import NodeClient
+from .tasks import TaskProgression, TaskProgressions
+
+__all__ = ["TaskManager"]
+
+
+class TaskManager(ContainerManager):
+    """
+    The Task Manager is responsible to distribute tasks depending on device resources.
+    It is composed with a Queue of Tasks.
+
+    Parameters
+    ----------
+    light_service_port : int
+        Light service port
+    data_folder : str
+        Data folder
+    """
+
+    tracer: Tracer
+    node_client: NodeClient
+
+    def __init__(
+        self,
+        light_service_port: int,
+    ) -> None:
+        # Initialize the TaskManager
+        self.ls_port = light_service_port
+        self.temp_folder = TemporaryDirectory()
+
+        # Temporary directory for the module python files
+        self.module_folder = Path(self.temp_folder.name)
+        self.tasks = TaskProgressions()
+
+        # Initialize the task queue
+        self.task_queue = deque()
+        self.future_task = None
+
+        ContainerManager.__init__(self)
+
+    async def check_tasks(self) -> List[TaskProgression]:
+        """
+        Check which tasks are active and update new status to the Manager
+
+        Returns
+        -------
+        List[TaskProgression]
+            List of updated tasks
+        """
+        if len(self.task_queue) > 0 or len(self.tasks._env_ids) > 0:
+            self.tracer.debug(
+                f"Checking tasks (task_queue={len(self.task_queue)}, tasks={len(self.tasks._env_ids)})"
+            )
+        await self.start_next_task()
+        updated = await self.update_active_tasks()
+        updated.extend(await self.update_finished_tasks())
+        return updated
+
+    # TODO : make subfunctions with decorator
+    async def start_next_task(self):
+        """
+        Start next start
+        """
+        if self.future_task is None and len(self.task_queue) > 0:
+            task = self.task_queue.pop()
+            self.future_task = asyncio.create_task(self.exec_task(task))
+        await self.check_future_task()
+
+    async def check_future_task(self):
+        """
+        Check future task
+        """
+        if self.future_task is None or not self.future_task.done():
+            return
+
+        try:
+            task_id = await self.future_task
+            self.tracer.info(f"Task {task_id} deployed from queue")
+        except MantaGRPCError as e:
+            self.tracer.exception(f"Error when executing task: {e}")
+        except MantaError as manta_error:
+            self.tracer.exception("Error when executing task.")
+            await self.node_client.send_error(manta_error)
+        except Exception as exc:
+            message = f"Error when executing task: {repr(exc)}"
+            self.tracer.exception(message)
+            await self.node_client.send_error(
+                MantaNodeError(message, metadata=self.tracer.metadata)
+            )
+
+        # Reinitialize the future task
+        self.future_task = None
+
+    # TODO : make subfunctions with decorator
+    async def update_active_tasks(self) -> List[TaskProgression]:
+        """
+        Update active tasks
+
+        Returns
+        -------
+        List[TaskProgression]
+            List of updated tasks
+        """
+        # Update the tasks status
+        updated = []
+        for task_progression in self.tasks.active_tasks():
+            try:
+                has_changed = await self.from_pending_to_running(task_progression)
+                if not task_progression.is_finished_task() and has_changed:
+                    updated.append(task_progression)
+            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)
+                )
+        return updated
+
+    # TODO : make subfunctions with decorator
+    async def update_finished_tasks(self) -> List[TaskProgression]:
+        """
+        Update finished tasks
+
+        Returns
+        -------
+        List[TaskProgression]
+            List of updated tasks
+        """
+        # Send logs of the finished tasks
+        updated = []
+        for task_progression in self.tasks.finished_tasks():
+            task_progression.task_status = TaskStatus.STOPPED
+            self.tracer.info(
+                f"Task {task_progression.task_id} finished: {task_progression.task_status}"
+            )
+            updated.append(task_progression)
+        return updated
+
+    async def from_pending_to_running(self, task_progression: TaskProgression) -> bool:
+        """
+        Change a task progression from pending status (supposed) to running status
+
+        Parameters
+        ----------
+        task_progression : TaskProgression
+            Task Progression
+
+        Returns
+        -------
+        bool
+            If task status has changed
+        """
+        task_id = task_progression.task_id
+        swarm_id = task_progression.swarm_id
+        old_task_status = task_progression.task_status
+        new_task_status = await self.get_task_status(swarm_id, task_id)
+        task_progression.task_status = new_task_status
+        return old_task_status != new_task_status
+
+    async def from_completed_to_pending(
+        self, task_progression: TaskProgression, task: MqttTask
+    ):
+        """
+        Change a task progression from completed status to pending status
+
+        Parameters
+        ----------
+        task_progression : TaskProgression
+            Task Progression
+        task : MqttTask
+            Task from MQTT
+        """
+        # Update the task information
+        task_progression.task_update.iteration = task.iteration
+        task_progression.task_update.circular = task.circular
+        task_progression.task_update.task_status = TaskStatus.PENDING
+
+        # Restart the container
+        start_time = await self.restart_container(
+            task_progression.swarm_id, task_progression.task_id
+        )
+        task_progression.start_time = start_time
+
+    async def create_task(self, task: MqttTask):
+        """
+        Create a task
+
+        Parameters
+        ----------
+        task : MqttTask
+           Task information from MQTT
+        """
+        task_id = ID(task.task_id)
+        swarm_id = ID(task.swarm_id)
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=task.iteration,
+            circular=task.circular,
+        )
+        self.tracer.debug(f"Creating task {task_id} for {swarm_id}")
+
+        # Create a temporary file and write the module code
+        # (even if the file does not exist, it will be created)
+        module_file = self.module_folder / f"{swarm_id.oid}_{task_id.oid}.pyz"
+        module_file.unlink(missing_ok=True)
+        module_file.touch(exist_ok=True)
+        module_file.write_bytes(task.payload)
+
+        container_file = "/usr/src/module.pyz"
+
+        network = None
+        if task.network != str(None):
+            if task.auth_key != str(None) or task.login_server != str(None):
+                self.tracer.debug(
+                    f"Task {task_id} has auth_key {task.auth_key} and login_server {task.login_server}"
+                )
+                network = await self.ensure_tailscale_service(
+                    swarm_id,
+                    task.auth_key,
+                    task.login_server,
+                    task.alias,
+                    self.module_folder,
+                )
+            else:
+                self.tracer.error(
+                    f"Task {task_id} has network {task.network} but no auth_key or login_server"
+                )
+                raise MantaNodeError(
+                    f"Task {task_id} has network {task.network} but no auth_key or login_server",
+                    metadata=self.tracer.metadata,
+                )
+
+        # Deploy the task
+        start_time = await self.run_container(
+            swarm_id=swarm_id,
+            task_id=task_id,
+            image=task.image,
+            command=f"python {container_file}",
+            volumes={module_file.absolute(): {"bind": container_file, "mode": "ro"}},
+            environment={
+                "TASK_ID": task_id.xid,
+                "SWARM_ID": swarm_id.xid,
+                "RPC_PORT": self.ls_port,
+                "RPC_HOST": "host.docker.internal",
+            },
+            gpu=task.gpu,
+            network=network,
+            alias=task.alias,
+        )
+
+        # Update the task information
+        self.tasks.append(
+            TaskProgression(
+                task_update=TaskUpdate(
+                    swarm_id=swarm_id.oid,
+                    node_id=self.tracer.metadata.node_id,
+                    task_id=task_id.oid,
+                    iteration=task.iteration,
+                    circular=task.circular,
+                    task_status=TaskStatus.PENDING,
+                ),
+                module_file=module_file,
+                start_time=start_time,
+            )
+        )
+
+    async def exec_task(self, task: MqttTask) -> ID:
+        """
+        Deploy a task
+
+        Parameters
+        ----------
+        task : MqttTask
+            Task
+
+        Returns
+        -------
+        ID
+            Task ID
+        """
+        task_id = ID(task.task_id)
+        swarm_id = ID(task.swarm_id)
+        self.tracer.info(f"Deploying task {task_id}")
+
+        if (swarm_id.tid, task_id.tid) in self.tasks:
+            await self.restart_task(self.tasks[(swarm_id.tid, task_id.tid)], task)
+        else:
+            await self.create_task(task)
+        return task_id
+
+    async def restart_task(self, task_progression: TaskProgression, task: MqttTask):
+        """
+        Restart a task
+
+        Parameters
+        ----------
+        task_progression : TaskProgression
+            Task Progression
+        task : MqttTask
+            Task from MQTT
+        """
+        swarm_id = task_progression.swarm_id
+        task_id = task_progression.task_id
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=task.iteration,
+            circular=task.circular,
+        )
+        # Use domain service to check if task should restart
+        if not TaskLifecycleService.should_restart_task(
+            task_progression.task_update.iteration,
+            task.iteration,
+            task_progression.task_update.circular,
+            task.circular,
+        ):
+            # Same iteration/circular - this is a duplicate message, ignore it
+            self.tracer.warning(
+                f"Ignoring duplicate deploy message for task {task_id} "
+                f"(iteration={task.iteration}, circular={task.circular}). "
+                f"Task already exists with same parameters."
+            )
+            return
+
+        self.tracer.debug(f"Task {task_id} is already created, restarting it")
+        await self.from_completed_to_pending(task_progression, task)
+        self.tracer.info(f"Task {task_id} restarted successfully")
+
+    async def get_task_logs(
+        self, task_progression: TaskProgression
+    ) -> Tuple[datetime, bytes]:
+        """
+        Send the logs to the Manager
+
+        Parameters
+        ----------
+        task_id : str
+            Task ID
+
+        Raises
+        ------
+        TaskManagerException
+            If an error occurs during the add_task_logs
+        """
+        task_id = task_progression.task_id
+        swarm_id = task_progression.swarm_id
+        self.tracer.set_metadata(
+            task_id=task_id,
+            swarm_id=swarm_id,
+            iteration=task_progression.task_update.iteration,
+            circular=task_progression.task_update.circular,
+        )
+        logs_bytes = await self.get_container_logs(
+            swarm_id,
+            task_id,
+            since=task_progression.start_time,
+        )
+
+        return self.tracer.store_raw_logs(logs_bytes)
+
+    async def deploy_task(self, message: bytes):
+        """
+        Convert the message into a MqttTask and add it to the
+        task queue
+
+        Parameters
+        ----------
+        message : bytes
+            Message
+
+        Raises
+        ------
+        MantaMQTTError
+            Could not convert the message
+        """
+        try:
+            task = MqttTask().parse(message)
+            task_id = ID(task.task_id)
+            swarm_id = ID(task.swarm_id)
+            self.tracer.set_metadata(
+                task_id=task_id,
+                swarm_id=swarm_id,
+                iteration=task.iteration,
+                circular=task.circular,
+            )
+            self.task_queue.append(task)
+            self.tracer.info(f"Task {task_id} added to queue")
+        except Exception as exc:
+            trace_message = f"Unable to convert message into 'MqttTask': {repr(exc)}"
+            self.tracer.exception(trace_message)
+            raise MantaMQTTError(trace_message, metadata=self.tracer.metadata)
+
+    async def stop_task(self, message: bytes) -> Optional[TaskUpdate]:
+        """
+        Stop a task
+
+        Parameters
+        ----------
+        message : bytes
+            Message from MQTT client
+
+        Returns
+        -------
+        Optional[TaskUpdate]
+            Task update
+        """
+        try:
+            env_id = EnvId().parse(message)
+            swarm_id = ID(env_id.swarm_id)
+            task_id = ID(env_id.task_id)
+            task_progression = self.tasks[(swarm_id.tid, task_id.tid)]
+            self.tracer.set_metadata(
+                task_id=task_id,
+                swarm_id=swarm_id,
+                iteration=task_progression.task_update.iteration,
+                circular=task_progression.task_update.circular,
+            )
+
+            # Stop the container
+            await self.stop_container(swarm_id, task_id)
+
+            # Set the new TaskStatus "Stopped"
+            task_progression.task_status = TaskStatus.STOPPED
+            self.tracer.info(f"Task {task_id} stopped successfully")
+            return task_progression.task_update
+        except MantaGRPCError as e:
+            self.tracer.exception(f"gRPCError when stopping task: {e}")
+        except MantaError as manta_error:
+            self.tracer.error("Unable to stop task.")
+            raise manta_error
+        except Exception as exc:
+            trace_message = f"Unable to stop task: {repr(exc)}"
+            self.tracer.exception(trace_message)
+            raise MantaNodeError(trace_message, metadata=self.tracer.metadata)
+
+    # TODO : add decorator for loop
+    async def stop_swarm(self, message: bytes):
+        """
+        Stop a swarm
+
+        Parameters
+        ----------
+        swarm_id : bytess
+            Swarm ID
+        """
+        swarm_id = ID(PairId().parse(message))
+        self.tracer.set_metadata(swarm_id=swarm_id)
+
+        def match_swarm_id(task_progression: TaskProgression):
+            return task_progression.swarm_id == swarm_id
+
+        # Convert to list to avoid concurrent modification during iteration
+        tasks_to_stop = list(self.tasks.filter(match_swarm_id))
+        for task_progression in tasks_to_stop:
+            try:
+                await self.stop_and_remove_task(task_progression)
+            except MantaGRPCError as e:
+                self.tracer.exception(f"gRPCError when stopping swarm: {e}")
+            except MantaError as manta_error:
+                self.tracer.exception(
+                    f"Unable to stop task {task_progression.task_id}"
+                    f" of swarm {swarm_id}"
+                )
+                await self.node_client.send_error(manta_error)
+            except Exception as exc:
+                trace_message = (
+                    f"Unable to stop task {task_progression.task_id}"
+                    f" of swarm {swarm_id}: {repr(exc)}"
+                )
+                self.tracer.exception(trace_message)
+                await self.node_client.send_error(
+                    MantaNodeError(trace_message, metadata=self.tracer.metadata)
+                )
+
+        self.tracer.info(f"All tasks in swarm {swarm_id} stopped successfully")
+
+    async def stop_and_remove_task(self, task_progression: TaskProgression):
+        """
+        Stop task's container, remove the task's container and remove the task
+
+        Parameters
+        ----------
+        task_progression : TaskProgression
+            Task Progression
+        """
+        await self.stop_container(
+            task_progression.swarm_id,
+            task_progression.task_id,
+        )
+        await self.remove_container(
+            task_progression.swarm_id,
+            task_progression.task_id,
+        )
+        self.tasks.pop(
+            (
+                task_progression.swarm_id.tid,
+                task_progression.task_id.tid,
+            )
+        )