control_plane.py

# =====================================================
# Apckeyl Framework
# Version 1.0
# control_plane.py
# =====================================================

"""
Apckeyl Control Plane.

Version 1.0

Центральный слой управления Apckeyl Framework.

На этом этапе Control Plane:

- принимает тип задачи;
- ищет подходящий Compute Module;
- возвращает информацию о модуле.

Сетевые запросы пока НЕ выполняются.
"""


from module_registry import registry


# =====================================================
# Control Plane
# =====================================================

class ControlPlane:

    def __init__(self, module_registry=None):

        if module_registry is None:

            module_registry = registry

        self.registry = module_registry


    # =================================================
    # Resolve Task
    # =================================================

    def resolve_task(self, task_type):

        if not task_type:

            raise ValueError(
                "task_type is required"
            )

        modules = self.registry.find_by_type(
            self._map_task_to_module_type(
                task_type
            )
        )

        if not modules:

            raise LookupError(
                f"No Compute Module available "
                f"for task type: {task_type}"
            )

        # -------------------------------------------------
        # First available module
        # -------------------------------------------------

        for module in modules:

            if module["status"] == "available":

                return module

        raise RuntimeError(
            f"No available Compute Module "
            f"for task type: {task_type}"
        )


    # =================================================
    # Task → Module Type
    # =================================================

    def _map_task_to_module_type(
        self,
        task_type
    ):

        task_mapping = {

            "image_upscale":
                "image_upscaler",

        }

        module_type = task_mapping.get(
            task_type
        )

        if module_type is None:

            raise LookupError(
                f"Unknown task type: {task_type}"
            )

        return module_type


    # =================================================
    # Get Module
    # =================================================

    def get_module(
        self,
        module_id
    ):

        module = self.registry.get(
            module_id
        )

        if module is None:

            raise LookupError(
                f"Unknown module: {module_id}"
            )

        return module


    # =================================================
    # List Modules
    # =================================================

    def list_modules(self):

        return self.registry.get_all()


# =====================================================
# Default Control Plane
# =====================================================

control_plane = ControlPlane()