diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 6e3fcdcf..5b61e5df 100644 --- a/docs/ARCHITECTURE.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index f68d2bf1..32450126 100644 --- a/pyproject.toml +++ b/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", diff --git a/sdk/src/flowmesh/models/workers.py b/sdk/src/flowmesh/models/workers.py index f576b93d..54cf5fad 100644 --- a/sdk/src/flowmesh/models/workers.py +++ b/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): diff --git a/src/server/dispatcher/base.py b/src/server/dispatcher/base.py index 4c2a189e..71a6e4e5 100644 --- a/src/server/dispatcher/base.py +++ b/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 = [] diff --git a/src/server/registries/worker.py b/src/server/registries/worker.py index f18b028a..03d7e5ff 100644 --- a/src/server/registries/worker.py +++ b/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: diff --git a/src/server/task/parser.py b/src/server/task/parser.py index d1671359..efb8b62d 100644 --- a/src/server/task/parser.py +++ b/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 diff --git a/src/shared/schemas/worker.py b/src/shared/schemas/worker.py index ca1a6010..62dca7c2 100644 --- a/src/shared/schemas/worker.py +++ b/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.", ) diff --git a/src/shared/tasks/components/model.py b/src/shared/tasks/components/model.py index 4b4d4f80..f060672e 100644 --- a/src/shared/tasks/components/model.py +++ b/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 diff --git a/src/shared/tasks/specs/__init__.py b/src/shared/tasks/specs/__init__.py index 895058a3..b9c0f269 100644 --- a/src/shared/tasks/specs/__init__.py +++ b/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", diff --git a/src/shared/tasks/specs/common.py b/src/shared/tasks/specs/common.py index e841e596..6949c128 100644 --- a/src/shared/tasks/specs/common.py +++ b/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 diff --git a/src/shared/tasks/specs/inference.py b/src/shared/tasks/specs/inference.py index 40746254..cd89727e 100644 --- a/src/shared/tasks/specs/inference.py +++ b/src/shared/tasks/specs/inference.py @@ -1,3 +1,4 @@ +from enum import StrEnum from typing import Literal from ..placeholders import TemplateBool, TemplateInt @@ -10,6 +11,14 @@ ) +class InferenceBackend(StrEnum): + """The executor backend an inference task resolves to.""" + + TRANSFORMERS = "transformers" + VLLM = "vllm" + AUTO = "auto" + + class InferenceSpecStrict(ModelInferSpecStrict): taskType: Literal[TaskType.INFERENCE] @@ -17,6 +26,12 @@ class InferenceSpecStrict(ModelInferSpecStrict): parallel: ParallelSpec | None = None enforce_cpu: bool | None = None + def backend(self) -> InferenceBackend: + return _inference_backend(self) + + def validate_dispatchable(self) -> None: + _validate_inference_dispatchable(self) + class InferenceSpecTemplate(ModelInferSpecTemplate): taskType: Literal[TaskType.INFERENCE] @@ -24,3 +39,65 @@ class InferenceSpecTemplate(ModelInferSpecTemplate): sloSeconds: TemplateInt | None = None parallel: ParallelSpecTemplate | None = None enforce_cpu: TemplateBool | None = None + + def backend(self) -> InferenceBackend: + return _inference_backend(self) + + def validate_dispatchable(self) -> None: + _validate_inference_dispatchable(self) + + +def _inference_backend( + spec: InferenceSpecStrict | InferenceSpecTemplate, +) -> InferenceBackend: + """Classify the executor backend an inference task resolves to. + + The HF transformers executor runs on a GPU when one is available, so only ``VLLM`` + (explicit vLLM config or adapters) strictly requires a GPU. ``AUTO`` is the unhinted + case: the runner prefers vLLM but falls back to the transformers executor when vLLM + is absent. + """ + if spec.enforce_cpu is True: + return InferenceBackend.TRANSFORMERS + model = spec.model + if model is None: + return InferenceBackend.AUTO + if model.vllm or model.adapters: + return InferenceBackend.VLLM + if model.transformers: + return InferenceBackend.TRANSFORMERS + return InferenceBackend.AUTO + + +def _validate_inference_dispatchable( + spec: InferenceSpecStrict | InferenceSpecTemplate, +) -> None: + model = spec.model + if model and (adapters := model.adapters): + for adapter in adapters: + if not adapter.path and not adapter.url and not adapter.task_id: + raise ValueError( + f"adapter {adapter.name or adapter.type!r} specifies no path, " + "url, or task_id and cannot be loaded." + ) + + if isinstance(spec.enforce_cpu, str): # Unresolved template placeholder + return + if spec.enforce_cpu is True: + if model and model.vllm: + raise ValueError( + "enforce_cpu is set but the model configures a vLLM backend." + ) + return + + if spec.backend() is not InferenceBackend.VLLM: + return + hardware = spec.resources.hardware if spec.resources else None + gpu = hardware.gpu if hardware else None + if gpu and gpu.count is not None and gpu.count >= 1: + return + raise ValueError( + "inference task resolves to the vLLM backend but requests no GPU. Set " + "spec.resources.hardware.gpu.count >= 1, or use enforce_cpu / a " + "transformers-only model for CPU inference." + ) diff --git a/src/worker/README.md b/src/worker/README.md index bac10252..efca71bf 100644 --- a/src/worker/README.md +++ b/src/worker/README.md @@ -119,10 +119,11 @@ The SSH session container is labelled with `flowmesh.ssh.task_id` and `flowmesh.ssh.worker_id` so the server can clean it up if the worker container is stopped externally. -At startup the worker initializes the SSH executor only if the Docker daemon is -reachable, and advertises its `ssh` capability accordingly; the dispatcher then -selects it for SSH tasks only when it can actually run them. (The supervisor -mounts the Docker socket only when SSH is enabled for the node; see +At startup the worker advertises the task types it can service, and the +dispatcher only routes a task to a worker that advertises its type. A worker +advertises a type only when its executor came up: SSH requires a reachable +Docker daemon, so `ssh` is offered only on Docker-capable workers. (The +supervisor mounts the Docker socket only when SSH is enabled for the node; see `ENABLE_SSH_BY_DEFAULT` below.) Relevant env vars for SSH tasks: diff --git a/src/worker/executors/agent_executor.py b/src/worker/executors/agent_executor.py index 87d52312..6e2f926c 100644 --- a/src/worker/executors/agent_executor.py +++ b/src/worker/executors/agent_executor.py @@ -19,6 +19,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import AgentSpecStrict +from shared.tasks.task_type import TaskType from .base_executor import ExecutionError, Executor, ExecutorTask from .utils.checkpoints import maybe_upload_artifacts, write_executor_result @@ -66,6 +67,7 @@ class AgentExecutor(Executor): """Agent executor using youtu-agent (utu) framework""" name = "agent" + supported_task_types = frozenset({TaskType.AGENT}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/api_executor.py b/src/worker/executors/api_executor.py index 327da4cc..5fe99e9e 100644 --- a/src/worker/executors/api_executor.py +++ b/src/worker/executors/api_executor.py @@ -9,6 +9,7 @@ from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import ApiSpecStrict +from shared.tasks.task_type import TaskType from .base_executor import ExecutionError, Executor, ExecutorTask @@ -40,6 +41,7 @@ class APIExecutor(Executor): """ name = "api" + supported_task_types = frozenset({TaskType.API}) # ---- Class-level connection pool (shared across all instances) ---- _clients: ClassVar[dict[_ClientKey, httpx.Client]] = {} diff --git a/src/worker/executors/base_executor.py b/src/worker/executors/base_executor.py index 5a57c53f..3b2b6e76 100644 --- a/src/worker/executors/base_executor.py +++ b/src/worker/executors/base_executor.py @@ -29,11 +29,12 @@ def run(self, task: ExecutorTask, out_dir: Path) -> MyResult: import json from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, TypeVar +from typing import Any, ClassVar, TypeVar from shared.schemas.result import BaseExecutorResult from shared.tasks import MergedChildTaskStrict from shared.tasks.specs import TaskSpecStrictBase +from shared.tasks.task_type import TaskType from shared.tasks.worker_message import WorkerHardware, WorkerTaskMessage from worker.config import WorkerConfig from worker.lifecycle import Lifecycle @@ -67,8 +68,10 @@ class Executor(ABC): Subclasses must implement `run` and may override `prepare` and `teardown`. """ - #: Human-readable identifier for logging/telemetry name: str = "executor" + """Human-readable identifier for logging/telemetry""" + supported_task_types: ClassVar[frozenset[TaskType]] = frozenset() + """Types of tasks this executor can service""" def __init__( self, @@ -172,6 +175,7 @@ class EchoResult(BaseExecutorResult): class EchoExecutor(Executor): name = "echo" + supported_task_types = frozenset({TaskType.ECHO}) def run(self, task: ExecutorTask, out_dir: Path) -> EchoResult: return EchoResult( diff --git a/src/worker/executors/data_profiling_executor.py b/src/worker/executors/data_profiling_executor.py index 57f4249c..50644b73 100644 --- a/src/worker/executors/data_profiling_executor.py +++ b/src/worker/executors/data_profiling_executor.py @@ -10,6 +10,7 @@ from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import DataProfilingSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.json import to_json_serializable, validate_keys from ..connectors import get_connector_from_spec @@ -31,6 +32,7 @@ class DataProfilingExecutor(DataMixin, Executor): """Executor that estimates SQL query costs by sampling SQL template params.""" name = "data_profiling" + supported_task_types = frozenset({TaskType.DATA_PROFILING}) def run(self, task: ExecutorTask, out_dir: Path) -> DataProfilingResult: spec = self.require_spec(task, DataProfilingSpecStrict) diff --git a/src/worker/executors/data_retrieval_executor.py b/src/worker/executors/data_retrieval_executor.py index 7e7eab3f..9847e5cd 100644 --- a/src/worker/executors/data_retrieval_executor.py +++ b/src/worker/executors/data_retrieval_executor.py @@ -19,6 +19,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import DataRetrievalSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.json import validate_keys from ..connectors import LumidDataConnector, PostgreSQLConnector, S3Connector @@ -40,6 +41,7 @@ class DataRetrievalResult(BaseExecutorResult): class DataRetrievalExecutor(DataMixin, Executor): name = "data_retrieval" + supported_task_types = frozenset({TaskType.DATA_RETRIEVAL}) def run(self, task: ExecutorTask, out_dir: Path) -> DataRetrievalResult: spec = self.require_spec(task, DataRetrievalSpecStrict) diff --git a/src/worker/executors/diffusers_executor.py b/src/worker/executors/diffusers_executor.py index 93b640a7..698e15cb 100644 --- a/src/worker/executors/diffusers_executor.py +++ b/src/worker/executors/diffusers_executor.py @@ -18,6 +18,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import DiffusionSpecStrict +from shared.tasks.task_type import TaskType from ..utils.logging import configure_hf_library_logging from .base_executor import ExecutionError, Executor, ExecutorTask @@ -54,6 +55,7 @@ class DiffusersExecutor(DataMixin, Executor): """Executor that runs text-to-image generation via Hugging Face Diffusers.""" name = "diffusers" + supported_task_types = frozenset({TaskType.DIFFUSION}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/dpo_executor.py b/src/worker/executors/dpo_executor.py index 2229059a..4354492e 100644 --- a/src/worker/executors/dpo_executor.py +++ b/src/worker/executors/dpo_executor.py @@ -28,6 +28,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import DPOSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.manifest import scratch_dir from ..utils.logging import configure_hf_library_logging @@ -62,6 +63,7 @@ class DPOExecutor(TrainingMixin, Executor): """DPO training executor using TRL library.""" name = "dpo_executor" + supported_task_types = frozenset({TaskType.DPO}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/echo_executor.py b/src/worker/executors/echo_executor.py index 1a03662f..fcf69562 100644 --- a/src/worker/executors/echo_executor.py +++ b/src/worker/executors/echo_executor.py @@ -4,6 +4,7 @@ from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import EchoSpecStrict +from shared.tasks.task_type import TaskType from .base_executor import ExecutionError, Executor, ExecutorTask from .mixins.data import DataMixin @@ -22,6 +23,7 @@ class EchoResult(BaseExecutorResult): class EchoExecutor(DataMixin, Executor): name = "echo" + supported_task_types = frozenset({TaskType.ECHO}) def _append_outputs(self, out_items: list[dict[str, Any]], value: Any) -> None: if isinstance(value, list): diff --git a/src/worker/executors/image_classification_executor.py b/src/worker/executors/image_classification_executor.py index 17729915..86a7e213 100644 --- a/src/worker/executors/image_classification_executor.py +++ b/src/worker/executors/image_classification_executor.py @@ -31,6 +31,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import ImageClassificationTrainingSpecStrict +from shared.tasks.task_type import TaskType from ..utils.logging import configure_hf_library_logging from .base_executor import ExecutionError, Executor, ExecutorTask @@ -64,6 +65,7 @@ class ImageClassificationTrainingResult(BaseExecutorResult): class ImageClassificationTrainingExecutor(TrainingMixin, Executor): name = "image_classification_training_executor" + supported_task_types = frozenset({TaskType.IMAGE_CLASSIFICATION_TRAINING}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/lora_sft_executor.py b/src/worker/executors/lora_sft_executor.py index 8c5a6365..6a8423f8 100644 --- a/src/worker/executors/lora_sft_executor.py +++ b/src/worker/executors/lora_sft_executor.py @@ -21,6 +21,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import LoRASFTSpecStrict +from shared.tasks.task_type import TaskType from ..utils.logging import configure_hf_library_logging from .base_executor import ExecutionError, Executor, ExecutorTask @@ -35,13 +36,17 @@ from .utils.huggingface import build_hf_load_kwargs, pick_torch_dtype try: - from peft import LoraConfig, PeftModel, TaskType, get_peft_model + from peft import LoraConfig, PeftModel + from peft import TaskType as PeftTaskType + from peft import get_peft_model except ImportError: if TYPE_CHECKING: - from peft import LoraConfig, PeftModel, TaskType, get_peft_model + from peft import LoraConfig, PeftModel + from peft import TaskType as PeftTaskType + from peft import get_peft_model else: LoraConfig = None - TaskType = None + PeftTaskType = None get_peft_model = None PeftModel = None @@ -64,6 +69,7 @@ class LoRASFTExecutor(TrainingMixin, Executor): """Execute LoRA-based supervised fine-tuning using TRL's SFTTrainer.""" name = "lora_sft_executor" + supported_task_types = frozenset({TaskType.LORA_SFT}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) @@ -80,7 +86,7 @@ def run(self, task: ExecutorTask, out_dir: Path) -> LoRAResult: if ( LoraConfig is None - or TaskType is None + or PeftTaskType is None or get_peft_model is None or PeftModel is None ): @@ -174,7 +180,7 @@ def run(self, task: ExecutorTask, out_dir: Path) -> LoRAResult: task_type_raw = str(lora_cfg.get("task_type", "CAUSAL_LM")).upper() try: - task_type = TaskType[task_type_raw] + task_type = PeftTaskType[task_type_raw] except KeyError as exc: raise ValueError( f"Unsupported LoRA task_type '{task_type_raw}'" diff --git a/src/worker/executors/omni_text2audio_executor.py b/src/worker/executors/omni_text2audio_executor.py index 780ed0ee..d1e4d9d4 100644 --- a/src/worker/executors/omni_text2audio_executor.py +++ b/src/worker/executors/omni_text2audio_executor.py @@ -54,6 +54,7 @@ from shared.schemas.governance import SpanType from shared.tasks.specs import TaskSpecStrictBase from shared.tasks.specs.omni import OmniText2AudioSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.parsing import to_float, to_int from .base_executor import ExecutionError, ExecutorTask @@ -77,6 +78,7 @@ class OmniText2AudioExecutor(OmniExecutorBase): """Generate background music with Omni diffusion sampling.""" name = EXECUTOR_NAME + supported_task_types = frozenset({TaskType.OMNI_TEXT2AUDIO}) _TASK_SPEC_TYPE = OmniText2AudioSpecStrict def prepare(self) -> None: diff --git a/src/worker/executors/omni_text2general_executor.py b/src/worker/executors/omni_text2general_executor.py index 3d03263c..04cde5c4 100644 --- a/src/worker/executors/omni_text2general_executor.py +++ b/src/worker/executors/omni_text2general_executor.py @@ -30,6 +30,7 @@ from shared.schemas.governance import SpanType from shared.tasks.specs import TaskSpecStrictBase from shared.tasks.specs.omni import OmniText2GeneralSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.parsing import as_list, to_bool, to_float, to_int, to_int_list from .base_executor import ExecutionError, ExecutorTask @@ -62,6 +63,7 @@ class OmniText2GeneralExecutor(OmniExecutorBase): """Generate narration/speech audio using Qwen3-Omni through vllm_omni.Omni.""" name = EXECUTOR_NAME + supported_task_types = frozenset({TaskType.OMNI_TEXT2GENERAL}) _TASK_SPEC_TYPE = OmniText2GeneralSpecStrict def prepare(self) -> None: diff --git a/src/worker/executors/omni_text2image_executor.py b/src/worker/executors/omni_text2image_executor.py index f11db25b..75f7de41 100644 --- a/src/worker/executors/omni_text2image_executor.py +++ b/src/worker/executors/omni_text2image_executor.py @@ -19,6 +19,7 @@ from shared.schemas.governance import SpanType from shared.tasks.specs import TaskSpecStrictBase from shared.tasks.specs.omni import OmniText2ImageSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.parsing import as_list from .base_executor import ExecutionError, ExecutorTask @@ -38,6 +39,7 @@ class OmniText2ImageExecutor(OmniExecutorBase): """Generate images using vllm_omni.Omni.""" name = EXECUTOR_NAME + supported_task_types = frozenset({TaskType.OMNI_TEXT2IMAGE}) _TASK_SPEC_TYPE = OmniText2ImageSpecStrict def prepare(self) -> None: diff --git a/src/worker/executors/omni_text2speech_executor.py b/src/worker/executors/omni_text2speech_executor.py index b37aa74a..911a1154 100644 --- a/src/worker/executors/omni_text2speech_executor.py +++ b/src/worker/executors/omni_text2speech_executor.py @@ -20,6 +20,7 @@ from shared.schemas.governance import SpanType from shared.tasks.specs import TaskSpecStrictBase from shared.tasks.specs.omni import OmniText2SpeechSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.parsing import as_list, to_int from .base_executor import ExecutionError, ExecutorTask @@ -47,6 +48,7 @@ class OmniText2SpeechExecutor(OmniExecutorBase): """Generate speech audio using vllm_omni.Omni.""" name = EXECUTOR_NAME + supported_task_types = frozenset({TaskType.OMNI_TEXT2SPEECH}) _TASK_SPEC_TYPE = OmniText2SpeechSpecStrict def prepare(self) -> None: diff --git a/src/worker/executors/ppo_executor.py b/src/worker/executors/ppo_executor.py index c475cb4f..b5950838 100644 --- a/src/worker/executors/ppo_executor.py +++ b/src/worker/executors/ppo_executor.py @@ -36,6 +36,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import PPOSpecStrict +from shared.tasks.task_type import TaskType from shared.utils.manifest import scratch_dir from shared.utils.parsing import safe_float, safe_int, to_bool @@ -418,6 +419,7 @@ class PPOExecutor(TrainingMixin, Executor): """PPO training executor using TRL library.""" name = "ppo_executor" + supported_task_types = frozenset({TaskType.PPO}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/rag_executor.py b/src/worker/executors/rag_executor.py index 06b3c705..ee80046a 100644 --- a/src/worker/executors/rag_executor.py +++ b/src/worker/executors/rag_executor.py @@ -18,6 +18,7 @@ from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import RagSpecStrict +from shared.tasks.task_type import TaskType from .base_executor import ExecutionError, Executor, ExecutorTask from .utils.graph_templates import Message, build_prompts_from_graph_template @@ -37,6 +38,7 @@ class RAGResult(BaseExecutorResult): class RAGExecutor(Executor): name = EXECUTOR_NAME + supported_task_types = frozenset({TaskType.RAG}) def run(self, task: ExecutorTask, out_dir: Path) -> RAGResult: start_ts = time.time() diff --git a/src/worker/executors/sft_executor.py b/src/worker/executors/sft_executor.py index 13be2586..b29c09d8 100644 --- a/src/worker/executors/sft_executor.py +++ b/src/worker/executors/sft_executor.py @@ -25,6 +25,7 @@ from shared.schemas.artifact import ArtifactRef from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import SFTSpecStrict, TaskSpecStrictBase +from shared.tasks.task_type import TaskType from shared.utils.manifest import scratch_dir from ..utils.logging import configure_hf_library_logging @@ -59,6 +60,7 @@ class SFTResult(BaseExecutorResult): class SFTExecutor(TrainingMixin, Executor): name = "sft_executor" + supported_task_types = frozenset({TaskType.SFT}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/ssh_executor.py b/src/worker/executors/ssh_executor.py index 3ad3e5ef..4ac97abf 100644 --- a/src/worker/executors/ssh_executor.py +++ b/src/worker/executors/ssh_executor.py @@ -37,6 +37,7 @@ SSHOutputSpec, SSHSpecStrict, ) +from shared.tasks.task_type import TaskType from shared.tasks.worker_message import WorkerHardware from shared.utils import new_ssh_session_id, parse_float_env, parse_mem_to_bytes from shared.utils.hardware import ( @@ -384,6 +385,7 @@ class SSHExecutor(Executor): """Executor for SSH tasks (interactive sessions and non-interactive jobs).""" name = "ssh" + supported_task_types = frozenset({TaskType.SSH}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/transformers_executor.py b/src/worker/executors/transformers_executor.py index 07b5d9da..7d5eba76 100644 --- a/src/worker/executors/transformers_executor.py +++ b/src/worker/executors/transformers_executor.py @@ -63,6 +63,7 @@ EmbeddingSpecStrict, InferenceSpecStrict, ) +from shared.tasks.task_type import TaskType from ..utils.logging import configure_hf_library_logging from .base_executor import ExecutionError, Executor, ExecutorTask @@ -127,6 +128,7 @@ class HFTransformersExecutor(InferenceMixin, Executor): """Executor that runs text generation via Hugging Face Transformers.""" name = "transformers" + supported_task_types = frozenset({TaskType.INFERENCE, TaskType.EMBEDDING}) def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) diff --git a/src/worker/executors/vllm_executor.py b/src/worker/executors/vllm_executor.py index bb94b21d..e2a7cf51 100644 --- a/src/worker/executors/vllm_executor.py +++ b/src/worker/executors/vllm_executor.py @@ -69,6 +69,7 @@ from shared.schemas.governance import SpanType from shared.schemas.result import BaseExecutorResult from shared.tasks.specs import InferenceSpecStrict +from shared.tasks.task_type import TaskType from .base_executor import ExecutionError, Executor, ExecutorTask from .mixins.data import InferenceEntry @@ -119,6 +120,7 @@ class VLLMExecutor(InferenceMixin, Executor): """Executor that runs text generation using vLLM based on a YAML spec.""" name = "vllm" + supported_task_types = frozenset({TaskType.INFERENCE}) summarization_template = """Summarize the following document concisely in 2-3 \ sentences. Focus on the main topic and key information. diff --git a/src/worker/main.py b/src/worker/main.py index 6a6ae2ff..52d2a8b5 100644 --- a/src/worker/main.py +++ b/src/worker/main.py @@ -3,6 +3,7 @@ import signal from shared.schemas.worker import WorkerCapabilities +from shared.tasks.task_type import TaskType from shared.tasks.worker_message import WorkerHardware from .config import WorkerConfig @@ -164,6 +165,17 @@ def init_executor(key: str, *, gpu_required: bool = False): return executors, default_executor +def build_capabilities( + executors: dict[str, Executor], + registry: dict[str, type[Executor] | None] | None = None, +) -> WorkerCapabilities: + registry = registry or EXECUTOR_REGISTRY + supported_task_types = frozenset[TaskType]().union( + *(cls.supported_task_types for key in executors if (cls := registry.get(key))) + ) + return WorkerCapabilities(supported_task_types=supported_task_types) + + def main() -> None: args = _parse_args() if args.collect_hw: @@ -209,9 +221,9 @@ def main() -> None: enable_mp_executors=cfg.enable_mp_executors, ) - capabilities = WorkerCapabilities(ssh="ssh" in executors) + capabilities = build_capabilities(executors) ssh_limits = cfg.ssh_limits - if capabilities.ssh: + if TaskType.SSH in capabilities.supported_task_types: if ssh_limits is None: logger.warning( "SSH resource cap not configured; SSH sessions will be able to access " diff --git a/src/worker/requirements/requirements.txt b/src/worker/requirements/requirements.txt index 3ec58ed9..d4b5a27e 100644 --- a/src/worker/requirements/requirements.txt +++ b/src/worker/requirements/requirements.txt @@ -12,7 +12,7 @@ boto3==1.41.5 chunkr-ai==0.3.7 colorama==0.4.6 colorlog==6.9.0 -crawl4ai==0.8.9 +crawl4ai==0.9.0 datasets==5.0.0 ddgs==9.10.0 diffusers==0.38.0 diff --git a/src/worker/runner.py b/src/worker/runner.py index f778fd49..9cc27b37 100644 --- a/src/worker/runner.py +++ b/src/worker/runner.py @@ -12,20 +12,13 @@ from shared.schemas.result import BaseExecutorResult from shared.tasks import MergedChildTaskStrict -from shared.tasks.specs import InferenceSpecStrict, TaskSpecStrictBase -from shared.tasks.worker_message import ( - HardwareUsage, - WorkerHardware, - WorkerTaskMessage, -) +from shared.tasks.specs import InferenceBackend, InferenceSpecStrict, TaskSpecStrictBase +from shared.tasks.worker_message import HardwareUsage, WorkerHardware, WorkerTaskMessage from shared.utils.manifest import prepare_output_dir, sync_manifest from shared.utils.time import now_iso from .executors.base_executor import ExecutionError, Executor, TaskCancelledError -from .executors.utils.checkpoints import ( - get_http_destination, - write_executor_result, -) +from .executors.utils.checkpoints import get_http_destination, write_executor_result from .lifecycle import Lifecycle from .utils.logging import TaskLogEmitter @@ -236,18 +229,10 @@ def _maybe_emit_http( ) def _select_inference_executor_key(self, spec: InferenceSpecStrict) -> str: - enforce_cpu = spec.enforce_cpu - if enforce_cpu: + if spec.backend() is InferenceBackend.TRANSFORMERS: return "default" - - model_cfg = spec.model - if model_cfg is not None and model_cfg.has_lora_adapter(): - return "vllm_lora" - adapters = model_cfg.adapters if model_cfg else None - if adapters: + if (model_cfg := spec.model) and model_cfg.adapters: return "vllm_lora" - if model_cfg is not None and model_cfg.transformers and not model_cfg.vllm: - return "default" return "vllm" def _maybe_expire_active_executor(self) -> None: diff --git a/tests/server/registries/test_worker_registry.py b/tests/server/registries/test_worker_registry.py index 85ab2e9c..07509a31 100644 --- a/tests/server/registries/test_worker_registry.py +++ b/tests/server/registries/test_worker_registry.py @@ -13,6 +13,7 @@ HardwareRequirements, ResourcesSpec, ) +from shared.tasks.task_type import TaskType from shared.tasks.worker_message import ( CPUInfo, GpuInfo, @@ -237,34 +238,50 @@ def test_no_ssh_cap_behaves_as_before(self) -> None: class TestCapabilitySatisfies: - def test_ssh_task_requires_ssh_capable_worker(self) -> None: - w = _worker(capabilities=WorkerCapabilities(ssh=False)) + def test_task_requires_advertised_task_type(self) -> None: + w = _worker( + capabilities=WorkerCapabilities( + supported_task_types=frozenset({TaskType.ECHO}) + ) + ) assert capability_satisfies(w, _ssh_task()) is False - def test_ssh_task_accepts_ssh_capable_worker(self) -> None: - w = _worker(capabilities=WorkerCapabilities(ssh=True)) + def test_task_accepts_worker_advertising_task_type(self) -> None: + w = _worker( + capabilities=WorkerCapabilities( + supported_task_types=frozenset({TaskType.SSH}) + ) + ) assert capability_satisfies(w, _ssh_task()) is True - def test_default_worker_is_not_ssh_capable(self) -> None: - # An unreporting worker parses with the strict default. + def test_default_worker_supports_nothing(self) -> None: + # An unreporting worker parses with the strict (empty) default. assert capability_satisfies(_worker(), _ssh_task()) is False + assert capability_satisfies(_worker(), _task(cpu=2)) is False - def test_non_ssh_task_ignores_capabilities(self) -> None: - w = _worker(capabilities=WorkerCapabilities(ssh=False)) + def test_gate_covers_all_task_types(self) -> None: + w = _worker( + capabilities=WorkerCapabilities( + supported_task_types=frozenset({TaskType.ECHO}) + ) + ) assert capability_satisfies(w, _task(cpu=2)) is True + assert capability_satisfies(w, _ssh_task()) is False class TestParseCapabilities: def test_capabilities_round_trip(self) -> None: raw = { "status": "IDLE", - "capabilities_json": WorkerCapabilities(ssh=True).model_dump_json(), + "capabilities_json": WorkerCapabilities( + supported_task_types=frozenset({TaskType.SSH, TaskType.ECHO}) + ).model_dump_json(), } w = _parse_worker_from_redis("w-1", raw) assert w is not None - assert w.capabilities.ssh is True + assert w.capabilities.supported_task_types == {TaskType.SSH, TaskType.ECHO} def test_missing_capabilities_defaults_strict(self) -> None: w = _parse_worker_from_redis("w-1", {"status": "IDLE"}) assert w is not None - assert w.capabilities.ssh is False + assert w.capabilities.supported_task_types == frozenset() diff --git a/tests/server/task/test_parser.py b/tests/server/task/test_parser.py index d82c77c6..bda48a46 100644 --- a/tests/server/task/test_parser.py +++ b/tests/server/task/test_parser.py @@ -5,6 +5,7 @@ import pytest from server.task.parser import _deep_merge, parse_workflow +from shared.tasks.specs import TaskSpecTemplateBase class TestParseWorkflowNative: @@ -168,3 +169,79 @@ def test_empty_src(self) -> None: dst = {"a": 1} result = _deep_merge(dst, {}) assert result == {"a": 1} + + +class TestInferenceSpecValidation: + def _parse_inference(self, body: str) -> TaskSpecTemplateBase: + doc = textwrap.dedent("""\ + apiVersion: flowmesh/v1 + kind: Task + metadata: + name: infer + spec: + taskType: inference + """) + textwrap.indent(textwrap.dedent(body), " ") + return parse_workflow(doc, format="native").tasks[0].task.spec + + def _gpu_count(self, spec: TaskSpecTemplateBase) -> int | None: + hardware = spec.resources.hardware if spec.resources else None + gpu = hardware.gpu if hardware else None + return gpu.count if gpu else None + + def test_vllm_without_gpu_is_rejected(self) -> None: + with pytest.raises(ValueError, match="vLLM backend but requests no GPU"): + self._parse_inference("model:\n vllm:\n gpu_memory_utilization: 0.9\n") + + def test_vllm_with_gpu_parses(self) -> None: + spec = self._parse_inference( + "model:\n vllm:\n gpu_memory_utilization: 0.9\n" + "resources:\n hardware:\n gpu:\n count: 1\n" + ) + assert self._gpu_count(spec) == 1 + + def test_non_vllm_backend_parses_without_gpu(self) -> None: + spec = self._parse_inference("data:\n items: ['x']\n") + assert self._gpu_count(spec) is None + + def test_placeholder_enforce_cpu_defers(self) -> None: + spec = self._parse_inference( + "enforce_cpu: ${inputs.cpu}\n" + "model:\n vllm:\n gpu_memory_utilization: 0.9\n" + ) + assert self._gpu_count(spec) is None + + def test_enforce_cpu_with_vllm_is_rejected(self) -> None: + with pytest.raises(ValueError, match="enforce_cpu is set but the model"): + self._parse_inference( + "enforce_cpu: true\nmodel:\n vllm:\n gpu_memory_utilization: 0.9\n" + ) + + def test_enforce_cpu_with_gpu_parses(self) -> None: + # enforce_cpu selects the HF transformers executor (not vLLM), which still + # runs on a GPU when one is available, so a GPU request is valid here. + spec = self._parse_inference( + "enforce_cpu: true\nresources:\n hardware:\n gpu:\n count: 1\n" + ) + assert self._gpu_count(spec) == 1 + + def test_enforce_cpu_alone_parses(self) -> None: + spec = self._parse_inference("enforce_cpu: true\n") + assert self._gpu_count(spec) is None + + def test_adapter_without_source_is_rejected(self) -> None: + with pytest.raises(ValueError, match="no path, url, or task_id"): + self._parse_inference("model:\n adapters:\n - type: lora\n") + + def test_adapter_with_path_parses(self) -> None: + spec = self._parse_inference( + "model:\n adapters:\n - type: lora\n path: /models/adapter\n" + "resources:\n hardware:\n gpu:\n count: 1\n" + ) + assert self._gpu_count(spec) == 1 + + def test_adapter_with_task_id_parses(self) -> None: + spec = self._parse_inference( + "model:\n adapters:\n - type: lora\n task_id: tsk-abc\n" + "resources:\n hardware:\n gpu:\n count: 1\n" + ) + assert self._gpu_count(spec) == 1 diff --git a/tests/shared/test_inference_spec.py b/tests/shared/test_inference_spec.py new file mode 100644 index 00000000..499c6b43 --- /dev/null +++ b/tests/shared/test_inference_spec.py @@ -0,0 +1,88 @@ +"""Tests for the inference spec: backend classification and dispatch validation.""" + +from typing import Any + +import pytest + +from shared.tasks.specs import InferenceBackend, InferenceSpecStrict + + +def _spec(**fields: Any) -> InferenceSpecStrict: + return InferenceSpecStrict.model_validate({"taskType": "inference", **fields}) + + +class TestInferenceBackend: + def test_enforce_cpu_is_transformers(self) -> None: + assert _spec(enforce_cpu=True).backend() is InferenceBackend.TRANSFORMERS + + def test_transformers_only_is_transformers(self) -> None: + spec = _spec(model={"transformers": {"torch_dtype": "float32"}}) + assert spec.backend() is InferenceBackend.TRANSFORMERS + + def test_explicit_vllm_is_vllm(self) -> None: + spec = _spec(model={"vllm": {"gpu_memory_utilization": 0.9}}) + assert spec.backend() is InferenceBackend.VLLM + + def test_transformers_and_vllm_is_vllm(self) -> None: + spec = _spec( + model={ + "transformers": {"torch_dtype": "float32"}, + "vllm": {"gpu_memory_utilization": 0.9}, + } + ) + assert spec.backend() is InferenceBackend.VLLM + + def test_empty_vllm_config_is_not_vllm(self) -> None: + # An empty vLLM dict is treated as unconfigured (matching the runner): with + # a transformers model it stays on transformers, otherwise it is AUTO. + with_transformers = _spec(model={"transformers": {"x": 1}, "vllm": {}}) + assert with_transformers.backend() is InferenceBackend.TRANSFORMERS + assert _spec(model={"vllm": {}}).backend() is InferenceBackend.AUTO + + def test_lora_adapter_is_vllm(self) -> None: + spec = _spec(model={"adapters": [{"type": "lora"}]}) + assert spec.backend() is InferenceBackend.VLLM + + def test_non_lora_adapter_is_vllm(self) -> None: + spec = _spec(model={"adapters": [{"type": "ia3"}]}) + assert spec.backend() is InferenceBackend.VLLM + + def test_no_model_is_auto(self) -> None: + assert _spec().backend() is InferenceBackend.AUTO + + def test_unhinted_model_is_auto(self) -> None: + spec = _spec(model={"source": {"identifier": "gpt2"}}) + assert spec.backend() is InferenceBackend.AUTO + + +class TestValidateDispatchable: + """Exercises the resolved-spec path used at dispatch (no placeholder deferral).""" + + def test_vllm_without_gpu_raises(self) -> None: + with pytest.raises(ValueError, match="requests no GPU"): + _spec( + model={"vllm": {"gpu_memory_utilization": 0.9}} + ).validate_dispatchable() + + def test_vllm_with_gpu_ok(self) -> None: + _spec( + model={"vllm": {"gpu_memory_utilization": 0.9}}, + resources={"hardware": {"gpu": {"count": 1}}}, + ).validate_dispatchable() + + def test_enforce_cpu_with_vllm_raises(self) -> None: + with pytest.raises(ValueError, match="enforce_cpu is set but the model"): + _spec(enforce_cpu=True, model={"vllm": {"x": 1}}).validate_dispatchable() + + def test_adapter_without_source_raises(self) -> None: + with pytest.raises(ValueError, match="no path, url, or task_id"): + _spec(model={"adapters": [{"type": "lora"}]}).validate_dispatchable() + + def test_adapter_with_task_id_ok(self) -> None: + _spec( + model={"adapters": [{"type": "lora", "task_id": "tsk-abc"}]}, + resources={"hardware": {"gpu": {"count": 1}}}, + ).validate_dispatchable() + + def test_unhinted_auto_ok(self) -> None: + _spec(model={"source": {"identifier": "gpt2"}}).validate_dispatchable() diff --git a/tests/worker/test_executor_registry.py b/tests/worker/test_executor_registry.py index b257bafe..361d28b2 100644 --- a/tests/worker/test_executor_registry.py +++ b/tests/worker/test_executor_registry.py @@ -1,11 +1,17 @@ """Tests for the executor registry and safe import mechanism.""" +from pathlib import Path from types import SimpleNamespace import pytest import worker.executors as executors_pkg +from shared.schemas.result import BaseExecutorResult +from shared.tasks.task_type import TaskType +from tests.worker.factories import make_worker_config from worker.executors import EXECUTOR_CLASS_NAMES, EXECUTOR_REGISTRY, IMPORT_ERRORS +from worker.executors.base_executor import Executor, ExecutorTask +from worker.main import build_capabilities class TestExecutorRegistry: @@ -79,3 +85,69 @@ def test_import_executor_rejects_non_executor( "not a subclass of Executor" in executors_pkg._IMPORT_ERRORS["NotAnExecutor"] ) + + +class _StubExecutor(Executor): + def run(self, task: ExecutorTask, out_dir: Path) -> BaseExecutorResult: + raise NotImplementedError + + +class _StubEcho(_StubExecutor): + supported_task_types = frozenset({TaskType.ECHO}) + + +class _StubInference(_StubExecutor): + supported_task_types = frozenset({TaskType.INFERENCE, TaskType.EMBEDDING}) + + +class TestSupportedTaskTypes: + def test_every_available_executor_declares_task_types(self) -> None: + for key, cls in EXECUTOR_REGISTRY.items(): + if cls is not None: + assert cls.supported_task_types, f"{key} declares no task types" + + def test_default_executor_serves_inference_and_embedding(self) -> None: + cls = EXECUTOR_REGISTRY["default"] + assert cls is not None + assert {TaskType.INFERENCE, TaskType.EMBEDDING} <= cls.supported_task_types + + +class _EmptyCaps(Executor): + """Stand-in for a live MPExecutor wrapper: reports no task types of its own.""" + + def run(self, task: ExecutorTask, out_dir: Path) -> BaseExecutorResult: + raise NotImplementedError + + +class TestBuildCapabilities: + @staticmethod + def _executors(*keys: str) -> dict[str, Executor]: + # Values report empty supported_task_types, mirroring MPExecutor wrappers + # which do not forward it; build_capabilities must read the registry class + # by key instead, or wrapped executors would advertise nothing (the + # inference/training regression). + inst = _EmptyCaps(make_worker_config()) + return {key: inst for key in keys} + + def test_unions_from_registry_class_not_instance(self) -> None: + registry: dict[str, type[Executor] | None] = { + "echo": _StubEcho, + "inference": _StubInference, + } + caps = build_capabilities(self._executors("echo", "inference"), registry) + assert caps.supported_task_types == { + TaskType.ECHO, + TaskType.INFERENCE, + TaskType.EMBEDDING, + } + + def test_skips_unknown_and_unavailable_keys(self) -> None: + registry: dict[str, type[Executor] | None] = { + "echo": _StubEcho, + "missing": None, + } + caps = build_capabilities(self._executors("echo", "missing", "ghost"), registry) + assert caps.supported_task_types == frozenset({TaskType.ECHO}) + + def test_empty_is_safe(self) -> None: + assert build_capabilities({}, {}).supported_task_types == frozenset() diff --git a/uv.lock b/uv.lock index e2465cb0..61981dc1 100644 --- a/uv.lock +++ b/uv.lock @@ -1047,7 +1047,7 @@ wheels = [ [[package]] name = "crawl4ai" -version = "0.8.9" +version = "0.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiofiles" }, @@ -1085,9 +1085,9 @@ dependencies = [ { name = "unclecode-litellm" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/90/2b/9e2b376ec8d4b803cc7b47b428a0ac762c59ce1215d0ab5a4aa9e7afc4a3/crawl4ai-0.8.9.tar.gz", hash = "sha256:00a3c97b9492a694f24fa66bc80b9f1533e5637745dc0e30914c88aac52763e3", size = 608119, upload-time = "2026-06-04T06:19:16.08Z" } +sdist = { url = "https://files.pythonhosted.org/packages/52/37/9b288598bb0049b918eeb00f029592c748465c8a3e819461e95d8b5f9dab/crawl4ai-0.9.0.tar.gz", hash = "sha256:00c1c3516d9ca24a9b5004523ae67974ace8a5baadc6e92e295afca73f077153", size = 591734, upload-time = "2026-06-18T09:31:08.674Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/97/6443bf99add1c16b16d992b8d0816a3ce7adecde37323d368bed36dd0aff/crawl4ai-0.8.9-py3-none-any.whl", hash = "sha256:7836cfc9d81ed7fc646d614e984402c336e5a682462f3d92ff37de121a4244d6", size = 516731, upload-time = "2026-06-04T06:19:14.348Z" }, + { url = "https://files.pythonhosted.org/packages/c9/19/2d634f26d9ad1281874c503c1043174042fef9af631dc10798f4f22c1544/crawl4ai-0.9.0-py3-none-any.whl", hash = "sha256:1296072d9cd92e79144c6823ed6f27ab3966f16ce44e79fcbe86943b3435c366", size = 499084, upload-time = "2026-06-18T09:31:07.27Z" }, ] [[package]] @@ -2178,7 +2178,7 @@ ci = [ { name = "codespell", specifier = ">=2.4.1" }, { name = "colorama", specifier = ">=0.4.6" }, { name = "colorlog", specifier = ">=6.9.0" }, - { name = "crawl4ai", specifier = ">=0.8.9" }, + { name = "crawl4ai", specifier = ">=0.9.0" }, { name = "cryptography", specifier = ">=48.0.1" }, { name = "datasets", specifier = ">=4.7.0" }, { name = "ddgs", specifier = ">=9.10.0" }, @@ -2296,7 +2296,7 @@ runtime-agent = [ { name = "arxiv", specifier = ">=2.2.0" }, { name = "beautifulsoup4", specifier = ">=4.14.3" }, { name = "chunkr-ai", specifier = ">=0.3.7" }, - { name = "crawl4ai", specifier = ">=0.8.9" }, + { name = "crawl4ai", specifier = ">=0.9.0" }, { name = "ddgs", specifier = ">=9.10.0" }, { name = "google-genai", specifier = ">=1.47.0" }, { name = "hydra-core", specifier = ">=1.3.2" }, @@ -2400,7 +2400,7 @@ runtime-worker-cpu = [ { name = "chunkr-ai", specifier = ">=0.3.7" }, { name = "colorama", specifier = ">=0.4.6" }, { name = "colorlog", specifier = ">=6.9.0" }, - { name = "crawl4ai", specifier = ">=0.8.9" }, + { name = "crawl4ai", specifier = ">=0.9.0" }, { name = "datasets", specifier = ">=4.7.0" }, { name = "ddgs", specifier = ">=9.10.0" }, { name = "diffusers", specifier = ">=0.38.0" },