"""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.")