gateway_routing.py

"""Deterministic, public-safe routing for private CPU and GPU backends.

This module deliberately uses only the Python standard library so it can be
shipped in the source-audited public gateway without pulling in ML code.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Literal, Mapping


CPU_PROFILE = "cpu"
GPU_PROFILE = "gpu"
DUAL_ROUTING_MODE = "dual"
SINGLE_GPU_ROUTING_MODE = "single-gpu"
VALID_PROFILES = frozenset({CPU_PROFILE, GPU_PROFILE})
VALID_ROUTING_MODES = frozenset({DUAL_ROUTING_MODE, SINGLE_GPU_ROUTING_MODE})

CPU_ONLY_BACKBONES = frozenset({"raw_pixels", "hog"})
DEEP_BACKBONES = frozenset(
    {
        "auto",
        "mobilenet_v3_small",
        "resnet18",
        "resnet50",
        "densenet121",
        "efficientnet_b0",
    }
)
GPU_REQUIRED_BACKENDS = frozenset({"automm", "yolo"})
GPU_REQUIRED_TASKS = frozenset(
    {"object_detection", "instance_segmentation", "semantic_segmentation"}
)
SUPPORTED_IMAGE_PREDICTORS = frozenset(
    {
        "image_predictor",
        "image_agent_backend",
        "automm_image_backend",
        "yolo_image_backend",
    }
)
GPU_IMAGE_PREDICTORS = frozenset(
    {"automm_image_backend", "yolo_image_backend"}
)

_TASK_ALIASES = {
    "": "classification",
    "class": "classification",
    "classification": "classification",
    "image_classification": "classification",
    "object": "object_detection",
    "object_detection": "object_detection",
    "detection": "object_detection",
    "detect": "object_detection",
    "instance": "instance_segmentation",
    "instance_seg": "instance_segmentation",
    "instance_segmentation": "instance_segmentation",
    "semantic": "semantic_segmentation",
    "semantic_seg": "semantic_segmentation",
    "semantic_segmentation": "semantic_segmentation",
}
_BACKEND_ALIASES = {
    "": "auto",
    "auto": "auto",
    "🤖_auto": "auto",
    "existing": "existing",
    "existing_classification_pipeline": "existing",
    "existing_repo_classifier": "existing",
    "medical_image_agent": "medical_agent",
    "medical_agent": "medical_agent",
    "image_agent": "medical_agent",
    "agent": "medical_agent",
    "autogluon_automm": "automm",
    "automm": "automm",
    "yolo": "yolo",
    "benchmark_backends": "benchmark",
    "benchmark": "benchmark",
}
_DEVICE_ALIASES = {
    "": "auto",
    "auto": "auto",
    "auto_gpu_if_available": "auto",
    "cpu": "cpu",
    "cuda": "cuda",
    "cuda_requires_gpu_hardware": "cuda",
    "gpu": "cuda",
}
_BACKBONE_ALIASES = {
    "": "auto",
    "auto": "auto",
    "🤖_auto": "auto",
    "raw": "raw_pixels",
    "raw_pixel": "raw_pixels",
    "raw_pixels": "raw_pixels",
    "hog": "hog",
    "mobilenet_v3": "mobilenet_v3_small",
    "mobilenet_v3_small": "mobilenet_v3_small",
    "resnet18": "resnet18",
    "resnet50": "resnet50",
    "densenet121": "densenet121",
    "efficientnet_b0": "efficientnet_b0",
}
_PRESET_ALIASES = {
    "": "auto",
    "auto": "auto",
    "auto_agent_decides": "auto",
    "agent_decides": "auto",
    "none": "auto",
    "lc25000_densenet121": "lc25000_densenet121",
    "lc25000_resnet50": "lc25000_resnet50",
    "breakhis_efficientnet_b0": "breakhis_efficientnet_b0",
    "ham10000_efficientnet_b0": "ham10000_efficientnet_b0",
    "ham10000_mobilenet_v3": "ham10000_mobilenet_v3",
    "pcam_resnet18": "pcam_resnet18",
}
_PRESET_BACKBONES = {
    "lc25000_densenet121": "densenet121",
    "lc25000_resnet50": "resnet50",
    "breakhis_efficientnet_b0": "efficientnet_b0",
    "ham10000_efficientnet_b0": "efficientnet_b0",
    "ham10000_mobilenet_v3": "mobilenet_v3_small",
    "pcam_resnet18": "resnet18",
}
_BENCHMARK_ALIASES = {
    "classification": {
        "existing": "existing_repo_classifier",
        "existing_repo_classifier": "existing_repo_classifier",
        "automm": "automm_classification",
        "automm_classification": "automm_classification",
        "yolo": "yolo_classification",
        "yolo_classification": "yolo_classification",
    },
    "object_detection": {
        "yolo": "yolo_detection",
        "yolo_detection": "yolo_detection",
        "automm": "automm_detection",
        "automm_detection": "automm_detection",
    },
    "instance_segmentation": {
        "yolo": "yolo_instance_segmentation",
        "yolo_instance_segmentation": "yolo_instance_segmentation",
        "automm": "automm_instance_segmentation",
        "automm_instance_segmentation": "automm_instance_segmentation",
    },
    "semantic_segmentation": {
        "yolo": "yolo_semantic_segmentation",
        "yolo_semantic_segmentation": "yolo_semantic_segmentation",
        "automm": "automm_semantic_segmentation",
        "automm_semantic_segmentation": "automm_semantic_segmentation",
    },
}


class RoutingError(ValueError):
    """Raised when routing metadata is unsupported, ambiguous, or forged."""


@dataclass(frozen=True)
class RouteDecision:
    profile: Literal["cpu", "gpu"]
    reason_code: str
    explanation: str


def _token(value: Any) -> str:
    return "_".join(
        str(value or "").strip().lower().replace("-", " ").split()
    )


def _normalize(value: Any, aliases: Mapping[str, str], field: str) -> str:
    token = _token(value)
    try:
        return aliases[token]
    except KeyError as exc:
        raise RoutingError(f"Unknown {field} value: {value!r}.") from exc


def _routing_mode(value: Any) -> str:
    mode = str(value or DUAL_ROUTING_MODE).strip().lower()
    if mode not in VALID_ROUTING_MODES:
        raise RoutingError(
            f"Unsupported routing mode {value!r}; expected dual or single-gpu."
        )
    return mode


def _normalized_image_fields(
    payload: Mapping[str, Any], *, uploaded: bool = False
) -> dict[str, Any]:
    task = _normalize(payload.get("image_task"), _TASK_ALIASES, "image task")
    backend_key = "uploaded_image_backend" if uploaded else "image_backend"
    backend = _normalize(payload.get(backend_key), _BACKEND_ALIASES, "image backend")
    device = _normalize(payload.get("image_device"), _DEVICE_ALIASES, "image device")
    backbone = _normalize(
        payload.get("image_backbone"), _BACKBONE_ALIASES, "image backbone"
    )
    preset = _normalize(
        payload.get("cancer_preset"), _PRESET_ALIASES, "cancer preset"
    )

    if uploaded:
        if task != "classification":
            raise RoutingError("Uploaded image training supports classification only.")
        if backend not in {"auto", "existing", "medical_agent"}:
            raise RoutingError(
                "Uploaded image training supports the existing or Medical image agent backend only."
            )

    effective_backend = backend
    if backend == "auto":
        if task == "classification":
            effective_backend = "existing"
        elif task == "semantic_segmentation":
            effective_backend = "automm"
        else:
            effective_backend = "yolo"

    if effective_backend in {"existing", "medical_agent"} and task != "classification":
        raise RoutingError(
            f"The {effective_backend.replace('_', ' ')} backend supports classification only."
        )
    if effective_backend == "yolo" and task == "classification":
        raise RoutingError("YOLO classification is not supported by this application.")
    if effective_backend == "yolo" and task == "semantic_segmentation":
        raise RoutingError("YOLO semantic segmentation is not supported by this application.")
    if effective_backend == "automm" and task == "instance_segmentation":
        raise RoutingError("AutoMM instance segmentation is not supported by this application.")

    candidates: tuple[str, ...] = ()
    if effective_backend == "benchmark":
        text = str(payload.get("image_benchmark_backends") or "").strip()
        if text:
            aliases = _BENCHMARK_ALIASES[task]
            normalized: list[str] = []
            for item in text.split(","):
                candidate = _token(item)
                if not candidate:
                    continue
                if candidate not in aliases:
                    raise RoutingError(
                        f"Unknown benchmark candidate {item.strip()!r} for {task}."
                    )
                normalized.append(aliases[candidate])
            if not normalized:
                raise RoutingError("Benchmark candidate selection is empty.")
            candidates = tuple(normalized)

    if effective_backend == "medical_agent" and preset != "auto":
        backbone = _PRESET_BACKBONES[preset]

    return {
        "task": task,
        "requested_backend": backend,
        "backend": effective_backend,
        "device": device,
        "backbone": backbone,
        "preset": preset,
        "candidates": candidates,
    }


def _route_image(fields: Mapping[str, Any], mode: str) -> RouteDecision:
    task = str(fields["task"])
    requested_backend = str(fields["requested_backend"])
    backend = str(fields["backend"])
    device = str(fields["device"])
    backbone = str(fields["backbone"])
    candidates = tuple(fields["candidates"])

    if mode == SINGLE_GPU_ROUTING_MODE:
        return RouteDecision(
            GPU_PROFILE,
            "single_gpu_rollback",
            "operational rollback mode sends this valid request to the private GPU backend",
        )

    if task in GPU_REQUIRED_TASKS and requested_backend == "auto":
        return RouteDecision(
            GPU_PROFILE,
            "gpu_task_required",
            f"{task.replace('_', ' ')} requires GPU execution",
        )
    if backend in GPU_REQUIRED_BACKENDS:
        return RouteDecision(
            GPU_PROFILE,
            "gpu_backend_required",
            f"{backend.upper() if backend == 'yolo' else 'AutoMM'} requires GPU execution",
        )
    if task in GPU_REQUIRED_TASKS:
        return RouteDecision(
            GPU_PROFILE,
            "gpu_task_required",
            f"{task.replace('_', ' ')} requires GPU execution",
        )
    if backend == "benchmark":
        if not candidates or any(
            "automm" in candidate or "yolo" in candidate
            for candidate in candidates
        ):
            return RouteDecision(
                GPU_PROFILE,
                "benchmark_gpu_candidate",
                "the benchmark includes a GPU-required backend",
            )
        if device != "cpu":
            return RouteDecision(
                GPU_PROFILE,
                "benchmark_auto_gpu",
                "the CPU-compatible benchmark was not explicitly assigned to CPU",
            )
        return RouteDecision(
            CPU_PROFILE,
            "benchmark_cpu_only",
            "all benchmark candidates are CPU-compatible and CPU was selected",
        )
    if backend == "medical_agent":
        if device == "cpu":
            return RouteDecision(
                CPU_PROFILE,
                "medical_agent_cpu",
                "CPU-selected frozen Medical Image Agent feature extraction",
            )
        return RouteDecision(
            GPU_PROFILE,
            "medical_agent_gpu",
            "Auto/CUDA Medical Image Agent feature extraction uses GPU",
        )
    if backend != "existing":
        raise RoutingError(f"Unsupported effective image backend: {backend!r}.")
    if backbone in CPU_ONLY_BACKBONES:
        return RouteDecision(
            CPU_PROFILE,
            "cpu_only_backbone",
            f"{backbone.replace('_', ' ')} feature extraction is CPU-only",
        )
    if backbone not in DEEP_BACKBONES:
        raise RoutingError(f"Unsupported normalized image backbone: {backbone!r}.")
    if device == "cpu":
        return RouteDecision(
            CPU_PROFILE,
            "explicit_cpu_deep_features",
            "CPU-selected frozen deep feature extraction",
        )
    return RouteDecision(
        GPU_PROFILE,
        "deep_features_auto" if device == "auto" else "deep_features_cuda",
        "Auto/CUDA frozen deep feature extraction uses GPU",
    )


def route_existing(
    payload: Mapping[str, Any],
    *,
    dataset_is_image: bool,
    routing_mode: str = DUAL_ROUTING_MODE,
) -> RouteDecision:
    """Route one existing-dataset request without contacting either backend."""

    mode = _routing_mode(routing_mode)
    if not dataset_is_image:
        if mode == SINGLE_GPU_ROUTING_MODE:
            return RouteDecision(
                GPU_PROFILE,
                "single_gpu_rollback",
                "operational rollback mode sends this valid request to the private GPU backend",
            )
        return RouteDecision(
            CPU_PROFILE, "tabular_workflow", "tabular workflows run on private CPU"
        )
    return _route_image(_normalized_image_fields(payload), mode)


def route_uploaded(
    payload: Mapping[str, Any], *, routing_mode: str = DUAL_ROUTING_MODE
) -> RouteDecision:
    """Route one uploaded-dataset request without opening the upload."""

    mode = _routing_mode(routing_mode)
    data_type = str(payload.get("data_type") or "").strip()
    if data_type == "Tabular CSV":
        if mode == SINGLE_GPU_ROUTING_MODE:
            return RouteDecision(
                GPU_PROFILE,
                "single_gpu_rollback",
                "operational rollback mode sends this valid request to the private GPU backend",
            )
        return RouteDecision(
            CPU_PROFILE, "tabular_workflow", "tabular workflows run on private CPU"
        )
    if data_type != "Image Dataset":
        raise RoutingError("data_type must be Tabular CSV or Image Dataset.")
    return _route_image(_normalized_image_fields(payload, uploaded=True), mode)


def route_prediction(
    payload: Mapping[str, Any],
    manifest: Mapping[str, Any],
    *,
    routing_mode: str = DUAL_ROUTING_MODE,
) -> RouteDecision:
    """Route package replay from manifest metadata before model loading."""

    mode = _routing_mode(routing_mode)
    predictor = str(manifest.get("predictor") or "").strip()
    raw_artifact_type = str(manifest.get("artifact_type") or "").strip().lower()
    artifact_type = raw_artifact_type
    if not artifact_type:
        artifact_type = "image" if predictor in SUPPORTED_IMAGE_PREDICTORS else "tabular"
    if artifact_type not in {"tabular", "image"}:
        raise RoutingError(
            f"Unsupported artifact_type in model manifest: {artifact_type!r}."
        )
    device = _normalize(payload.get("image_device"), _DEVICE_ALIASES, "image device")

    if artifact_type == "tabular":
        if mode == SINGLE_GPU_ROUTING_MODE:
            return RouteDecision(
                GPU_PROFILE,
                "single_gpu_rollback",
                "operational rollback mode sends this valid request to the private GPU backend",
            )
        return RouteDecision(
            CPU_PROFILE, "tabular_package", "tabular package replay runs on private CPU"
        )

    if predictor and predictor not in SUPPORTED_IMAGE_PREDICTORS:
        raise RoutingError(f"Unsupported image predictor in model manifest: {predictor!r}.")
    if mode == SINGLE_GPU_ROUTING_MODE:
        return RouteDecision(
            GPU_PROFILE,
            "single_gpu_rollback",
            "operational rollback mode sends this valid request to the private GPU backend",
        )
    if predictor in GPU_IMAGE_PREDICTORS:
        return RouteDecision(
            GPU_PROFILE,
            "gpu_package_required",
            "this model package requires the private GPU backend",
        )
    if device == "cuda":
        return RouteDecision(
            GPU_PROFILE,
            "image_package_cuda",
            "CUDA-selected image package replay uses GPU",
        )
    return RouteDecision(
        CPU_PROFILE,
        "image_package_cpu",
        "CPU-compatible image package replay runs on private CPU",
    )


def declare_route(
    payload: Mapping[str, Any], decision: RouteDecision, routing_mode: str
) -> dict[str, Any]:
    """Copy a request and attach its immutable routing declaration."""

    mode = _routing_mode(routing_mode)
    result = dict(payload)
    result["execution_profile"] = decision.profile
    result["routing_mode"] = mode
    result["route_reason_code"] = decision.reason_code
    return result


def verify_declared_route(
    payload: Mapping[str, Any], decision: RouteDecision
) -> None:
    """Reject any request whose signed-in gateway declaration was altered."""

    mode = str(payload.get("routing_mode") or "").strip().lower()
    if mode not in VALID_ROUTING_MODES:
        raise RoutingError("Route declaration contains an unsupported routing mode.")
    expected = {
        "execution_profile": decision.profile,
        "route_reason_code": decision.reason_code,
    }
    for key, value in expected.items():
        if str(payload.get(key) or "") != value:
            raise RoutingError("Route declaration does not match the computed route.")
    if mode == SINGLE_GPU_ROUTING_MODE and decision.reason_code != "single_gpu_rollback":
        raise RoutingError("Route declaration does not match rollback routing mode.")
    if mode == DUAL_ROUTING_MODE and decision.reason_code == "single_gpu_rollback":
        raise RoutingError("Route declaration does not match dual routing mode.")