104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""Task dispatcher: composes TaskRunner / WorkflowRunner to execute assignments.
|
|
|
|
Capability: ``task-scheduler`` (closing the loop from ``assigned`` to ``executed``).
|
|
|
|
Composes the existing single-process execution entry points (``runtime.task.TaskRunner``,
|
|
``workflow.runner.WorkflowRunner``) without editing them, strictly through their
|
|
public ``run(task) -> Task`` / ``run(definition, device_id) -> WorkflowRun`` contracts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
from core.models import Task
|
|
|
|
if TYPE_CHECKING:
|
|
from cloud.store import CloudStore
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Assignment:
|
|
"""A scheduler-produced binding of a queued task to a specific device+host."""
|
|
|
|
task_id: str
|
|
device_id: str
|
|
host_id: str
|
|
goal: str | None
|
|
workflow_definition_id: str | None
|
|
|
|
|
|
class RemoteDispatchNotSupportedError(RuntimeError):
|
|
"""Raised when an assignment targets a host other than the local process."""
|
|
|
|
|
|
class UnknownWorkflowDefinitionError(RuntimeError):
|
|
"""Raised when a workflow-based assignment references an unknown definition id."""
|
|
|
|
|
|
# Callable type aliases; kept lazy so importing this module never imports runtime/workflow.
|
|
TaskRunnerFactory = Callable[[], "object"]
|
|
WorkflowRunnerFactory = Callable[[], "object"]
|
|
|
|
|
|
class TaskDispatcher:
|
|
"""Executes assignments via the existing task / workflow runners.
|
|
|
|
Local-only in this change: an assignment whose ``host_id`` does not match
|
|
the dispatcher's own ``local_host_id`` raises ``RemoteDispatchNotSupportedError``
|
|
rather than attempting any execution over the network.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
local_host_id: str,
|
|
task_runner_factory: TaskRunnerFactory,
|
|
workflow_runner_factory: WorkflowRunnerFactory,
|
|
store: "CloudStore",
|
|
) -> None:
|
|
self.local_host_id = local_host_id
|
|
self.task_runner_factory = task_runner_factory
|
|
self.workflow_runner_factory = workflow_runner_factory
|
|
self.store = store
|
|
|
|
def dispatch(self, assignment: Assignment) -> None:
|
|
if assignment.host_id != self.local_host_id:
|
|
raise RemoteDispatchNotSupportedError(
|
|
f"assignment {assignment.task_id} targets host "
|
|
f"{assignment.host_id!r}, but this dispatcher owns {self.local_host_id!r}"
|
|
)
|
|
|
|
if assignment.workflow_definition_id:
|
|
status = self._dispatch_workflow(assignment)
|
|
else:
|
|
status = self._dispatch_goal(assignment)
|
|
|
|
self.store.update_task(assignment.task_id, status=status)
|
|
|
|
def _dispatch_goal(self, assignment: Assignment) -> str:
|
|
task = Task(
|
|
goal=assignment.goal or "",
|
|
device_id=assignment.device_id,
|
|
)
|
|
result = self.task_runner_factory().run(task) # type: ignore[attr-defined]
|
|
result_status = getattr(result, "status", None)
|
|
return "done" if result_status == "completed" else "failed"
|
|
|
|
def _dispatch_workflow(self, assignment: Assignment) -> str:
|
|
runner = self.workflow_runner_factory()
|
|
store = getattr(runner, "store", None)
|
|
if store is None:
|
|
raise UnknownWorkflowDefinitionError(
|
|
"workflow runner has no store to load definitions from"
|
|
)
|
|
definition = store.get_definition(assignment.workflow_definition_id)
|
|
if definition is None:
|
|
raise UnknownWorkflowDefinitionError(
|
|
f"unknown workflow definition {assignment.workflow_definition_id!r}"
|
|
)
|
|
run = runner.run(definition, device_id=assignment.device_id) # type: ignore[attr-defined]
|
|
run_status = getattr(run, "status", None)
|
|
return "done" if run_status == "completed" else "failed"
|