agentic-mobile-control/packages/cloud-platform/cloud/scheduler.py

"""Task scheduler: bounded queue + pluggable assignment strategy.

Capability: ``task-scheduler``.

Mirrors the "string key -> swappable implementation" registry shape already used
by ``driver/registry.py``, ``perception/provider.py``, and ``workflow/conditions.py``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
from uuid import uuid4

from core.models import utc_now

if TYPE_CHECKING:
    from cloud.config import CloudConfig
    from cloud.dispatch import Assignment
    from cloud.pool import DevicePool, PooledDevice
    from cloud.store import CloudStore


ScheduledTaskStatus = Literal[
    "queued", "assigned", "dispatched", "done", "failed", "cancelled"
]


@dataclass(frozen=True)
class TaskConstraints:
    """Optional device constraints attached to a task submission."""

    driver_type: str | None = None
    capability_tags: list[str] = field(default_factory=list)
    target_host_id: str | None = None
    target_device_id: str | None = None


@dataclass
class ScheduledTask:
    """A task submitted to the cloud scheduler, awaiting or undergoing assignment."""

    id: str
    goal: str | None
    workflow_definition_id: str | None
    constraints: TaskConstraints
    status: ScheduledTaskStatus = "queued"
    assigned_device_id: str | None = None
    assigned_host_id: str | None = None
    attempt_count: int = 0
    lease_id: str | None = None
    lease_expires_at: datetime | None = None
    terminal_result: dict[str, Any] | None = None
    failure_reason: str | None = None
    updated_at: datetime | None = None
    created_at: datetime = field(default_factory=utc_now)
    progress_step_index: int | None = None
    progress_step_status: str | None = None
    progress_summary: str | None = None
    progress_updated_at: datetime | None = None
    cancel_requested_at: datetime | None = None


@runtime_checkable
class AssignmentStrategy(Protocol):
    """Selects one device from a pre-filterd list of matching candidates."""

    def select(
        self,
        task: ScheduledTask,
        candidates: "list[PooledDevice]",
    ) -> "PooledDevice | None": ...


class FifoMatchStrategy:
    """Default strategy: return the first candidate, or None if empty.

    The caller is responsible for pre-filtering to idle, constraint-matching
    devices in submission order. This keeps the strategy trivially replaceable.
    """

    def select(
        self,
        task: ScheduledTask,
        candidates: "list[PooledDevice]",
    ) -> "PooledDevice | None":
        if not candidates:
            return None
        return candidates[0]


FIFO_MATCH_STRATEGY_NAME = "fifo_match"
DEFAULT_STRATEGIES: dict[str, AssignmentStrategy] = {
    FIFO_MATCH_STRATEGY_NAME: FifoMatchStrategy(),
}


class UnknownAssignmentStrategyError(ValueError):
    """Raised when the configured default strategy is not in the registry."""


class QueueFullError(RuntimeError):
    """Raised when a submission would exceed ``config.max_queue_depth``."""


class TaskSubmissionValidationError(ValueError):
    """Raised when a submission has neither a goal nor a workflow_definition_id."""


class TaskScheduler:
    """Accepts task submissions and assigns queued tasks to idle devices."""

    def __init__(
        self,
        pool: "DevicePool",
        store: "CloudStore",
        config: "CloudConfig",
        *,
        strategies: dict[str, AssignmentStrategy] | None = None,
    ) -> None:
        self.pool = pool
        self.store = store
        self.config = config
        self._strategies = (
            dict(strategies) if strategies is not None else dict(DEFAULT_STRATEGIES)
        )
        if config.default_assignment_strategy not in self._strategies:
            raise UnknownAssignmentStrategyError(
                f"unknown assignment strategy {config.default_assignment_strategy!r}; "
                f"registered strategies: {sorted(self._strategies)}"
            )

    def submit(
        self,
        goal: str | None = None,
        workflow_definition_id: str | None = None,
        constraints: TaskConstraints | None = None,
    ) -> str:
        if not goal and not workflow_definition_id:
            raise TaskSubmissionValidationError(
                "a submission must specify either a goal or a workflow_definition_id"
            )
        if self.store.count_queued_tasks() >= self.config.max_queue_depth:
            raise QueueFullError(
                f"task queue is full ({self.config.max_queue_depth} queued)"
            )
        task = ScheduledTask(
            id=uuid4().hex,
            goal=goal,
            workflow_definition_id=workflow_definition_id,
            constraints=constraints or TaskConstraints(),
            status="queued",
            created_at=utc_now(),
        )
        self.store.enqueue_task(task)
        return task.id

    def assign(self) -> "list[Assignment]":
        from cloud.dispatch import Assignment

        strategy = self._strategies[self.config.default_assignment_strategy]
        assignments: list[Assignment] = []
        queued = self.store.list_queued_tasks()  # ordered oldest-first
        if not queued:
            return assignments

        devices = self.pool.list_devices()
        now = utc_now()
        reserved_device_ids = self.store.list_reserved_device_ids(now=now)
        for task in queued:
            candidates = [
                device
                for device in devices
                if device.device_id not in reserved_device_ids
                and device.status == "idle"
                and _matches(device, task.constraints)
            ]
            selected = strategy.select(task, candidates)
            if selected is None:
                continue
            leased = self.store.assign_task(
                task_id=task.id,
                host_id=selected.host_id,
                device_id=selected.device_id,
                lease_id=uuid4().hex,
                lease_expires_at=now
                + timedelta(seconds=self.config.lease_duration_seconds),
                now=now,
            )
            if leased is None:
                continue
            reserved_device_ids.add(selected.device_id)
            assignments.append(
                Assignment(
                    task_id=leased.task_id,
                    device_id=leased.device_id,
                    host_id=leased.host_id,
                    goal=leased.goal,
                    workflow_definition_id=leased.workflow_definition_id,
                )
            )
        return assignments


def _matches(device: "PooledDevice", constraints: TaskConstraints) -> bool:
    if device.mcp_busy:
        return False
    if constraints.target_host_id and device.host_id != constraints.target_host_id:
        return False
    if (
        constraints.target_device_id
        and device.device_id != constraints.target_device_id
    ):
        return False
    if constraints.driver_type and device.driver_type != constraints.driver_type:
        return False
    if constraints.capability_tags:
        device_tags = set(device.capability_tags)
        if not all(tag in device_tags for tag in constraints.capability_tags):
            return False
    return True