feat: extend capability gate to all task types (#76)

Author: kaiitunnz

docs/ARCHITECTURE.md

@@ -127,12 +127,12 @@ scripts/dev/            compile_protos, sync_requirements, check_env_examples
   `WorkerHardware`. The dispatcher's `_cached_worker_candidates` filters
   to workers whose cache covers the task's references; entries older
   than `WORKER_CACHE_TTL_SEC` are ignored.
-- **Worker capabilities.** Beyond hardware fit, each worker advertises a
-  `WorkerCapabilities` describing the task kinds it can service. The dispatcher
-  routes a task that requires a capability only to workers advertising it, and
-  excludes a worker that advertises none from those tasks rather than failing
-  them on it. SSH is one such capability — a worker advertises it when it can
-  launch session containers.
+- **Worker capabilities.** Beyond hardware fit, each worker advertises the set
+  of task types it can service, and the dispatcher routes a task only to workers
+  that advertise its type. A worker advertises a type only when its executor came
+  up — e.g. SSH requires a reachable Docker daemon, and training or omni types
+  require their (often GPU-only) dependencies — so a worker missing that executor
+  isn't a candidate, rather than being handed a task it would fail.
 - **Cursor pagination.** List endpoints accept `limit` and `before` /
   `after` cursors. The cursor is an opaque base64 of `(timestamp, id)`;
   do not parse client-side.

pyproject.toml

@@ -83,7 +83,7 @@ runtime-agent = [
     "arxiv>=2.2.0",
     "beautifulsoup4>=4.14.3",
     "chunkr-ai>=0.3.7",
-    "crawl4ai>=0.8.9",
+    "crawl4ai>=0.9.0",
     "ddgs>=9.10.0",
     "google-genai>=1.47.0",
     "hydra-core>=1.3.2",

sdk/src/flowmesh/models/workers.py

@@ -4,6 +4,8 @@
 
 from pydantic import BaseModel, Field, field_validator
 
+from .common import TaskType
+
 
 class CPUInfo(BaseModel):
     logical_cores: int | None = None
@@ -80,7 +82,7 @@ class SSHLimits(BaseModel):
 
 
 class WorkerCapabilities(BaseModel):
-    ssh: bool = False
+    supported_task_types: frozenset[TaskType] = Field(default_factory=frozenset)
 
 
 class Worker(BaseModel):

src/server/dispatcher/base.py

@@ -303,7 +303,7 @@ def dispatch_once(self, task_id: str) -> bool:
             self._requeue_task(task_id, reason="no_selection")
             return False
 
-        # 5.5. Plan task merge: coalesce sibling merge candidates onto this worker
+        # Plan task merge: coalesce sibling merge candidates onto this worker
         merged_children: list[str] = []
         if self._task_merge_enabled and self._task_merge_max_batch_size > 1:
             merged_children = self._runtime.plan_merge(
@@ -341,11 +341,21 @@ def dispatch_once(self, task_id: str) -> bool:
             self._fail_task(task_id, str(exc), payload={"error": str(exc)})
             return True
 
-        # 6.5. Conditional execution: skip dispatch if condition not met
+        # Re-validate resolved task spec
+        try:
+            rendered_task.spec.validate_dispatchable()
+        except ValueError as exc:
+            self._runtime.release_merge(task_id)
+            self._fail_task(
+                task_id, "spec_validation_failed", payload={"error": str(exc)}
+            )
+            return True
+
+        # Conditional execution: skip dispatch if condition not met
         if self._evaluate_condition_skip(task_id, rendered_task, record):
             return True
 
-        # 6.6. Resolve merged-child rendered payloads
+        # Resolve merged-child rendered payloads
         rendered_children: list[MergedChildTaskStrict] | None = None
         if record.merged_children:
             rendered_children = []

src/server/registries/worker.py

@@ -550,9 +550,7 @@ def hw_satisfies(worker: Worker, task: TaskEnvelope) -> bool:
 
 
 def capability_satisfies(worker: Worker, task: TaskEnvelope) -> bool:
-    if isinstance(task.spec, (SSHSpecStrict, SSHSpecTemplate)):
-        return worker.capabilities.ssh
-    return True
+    return task.spec.taskType in worker.capabilities.supported_task_types
 
 
 def _gpu_meets_requirements(hw: WorkerHardware, gpu_req: GPURequirements) -> bool:

src/server/task/parser.py

@@ -161,6 +161,10 @@ def _build_task_template(
             ) from exc
         raise ValueError(f"Invalid task payload{context}: {exc}") from exc
     _validate_ssh_access_mode(task, context)
+    try:
+        task.spec.validate_dispatchable()
+    except ValueError as exc:
+        raise ValueError(f"Invalid task payload{context}: {exc}") from exc
     return task
 
 

src/shared/schemas/worker.py

@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel, Field
 
+from shared.tasks.task_type import TaskType
+
 
 class WorkerStatus(StrEnum):
     UNKNOWN = "UNKNOWN"
@@ -33,8 +35,9 @@ class SSHLimits(BaseModel):
 class WorkerCapabilities(BaseModel):
     """Task capabilities a worker advertises to the dispatcher."""
 
-    ssh: bool = Field(
-        default=False, description="Whether the worker can run SSH session tasks."
+    supported_task_types: frozenset[TaskType] = Field(
+        default_factory=frozenset,
+        description="Types of tasks this worker can service.",
     )
 
 

src/shared/tasks/components/model.py

@@ -76,13 +76,6 @@ class ModelConfig(StrictBaseModel):
     diffusers: dict[str, Any] | None = None
     adapters: list[AdapterConfig] | None = None
 
-    def has_lora_adapter(self) -> bool:
-        adapters = self.adapters or []
-        for adapter in adapters:
-            if (adapter.type or "").strip().lower() == "lora":
-                return True
-        return False
-
 
 class ModelConfigTemplate(TemplateBaseModel):
     source: ModelSourceTemplate | None = None
@@ -91,10 +84,3 @@ class ModelConfigTemplate(TemplateBaseModel):
     transformers: dict[str, Any] | None = None
     diffusers: dict[str, Any] | None = None
     adapters: list[AdapterConfigTemplate] | None = None
-
-    def has_lora_adapter(self) -> bool:
-        adapters = self.adapters or []
-        for adapter in adapters:
-            if (adapter.type or "").strip().lower() == "lora":
-                return True
-        return False

src/shared/tasks/specs/__init__.py

@@ -5,7 +5,7 @@
     TaskSpecTemplateBase,
 )
 from .diffusion import DiffusionSpecStrict, DiffusionSpecTemplate
-from .inference import InferenceSpecStrict, InferenceSpecTemplate
+from .inference import InferenceBackend, InferenceSpecStrict, InferenceSpecTemplate
 from .misc import (
     AgentSpecStrict,
     AgentSpecTemplate,
@@ -56,6 +56,7 @@
     "DiffusionSpecTemplate",
     "ImageClassificationTrainingSpecStrict",
     "ImageClassificationTrainingSpecTemplate",
+    "InferenceBackend",
     "InferenceSpecStrict",
     "InferenceSpecTemplate",
     "LoRASFTSpecStrict",

src/shared/tasks/specs/common.py

@@ -91,6 +91,14 @@ def get_artifacts(self) -> list[str]:
             return []
         return artifacts.copy()
 
+    def validate_dispatchable(self) -> None:
+        """Validate spec-internal invariants for a runnable task.
+
+        Called at submit and again before dispatch. Overrides must raise ``ValueError``
+        for misconfigurations.
+        """
+        return None
+
 
 class TaskSpecTemplateBase(TemplateBaseModel):
     resources: ResourcesSpec | None = None
@@ -117,6 +125,15 @@ def get_artifacts(self) -> list[str]:
             return []
         return artifacts.copy()
 
+    def validate_dispatchable(self) -> None:
+        """Validate spec-internal invariants for a runnable task.
+
+        Called at submit and again before dispatch. Overrides must defer
+        placeholder-dependent checks and raise ``ValueError`` for genuine
+        misconfigurations.
+        """
+        return None
+
 
 type TaskSpecBase = TaskSpecStrictBase | TaskSpecTemplateBase