From 68a187c0d43f6c3fb2d518273eb18e310b0320fc Mon Sep 17 00:00:00 2001 From: Kristian Bendiksen Date: Tue, 18 Aug 2026 15:16:50 +0200 Subject: [PATCH 1/2] Add entry-point discovery for tasks and workflows --- README.md | 32 ++++++++++- pyproject.toml | 2 +- taskmaestro/__init__.py | 22 +++++++- taskmaestro/discovery.py | 96 ++++++++++++++++++++++++++++++++ taskmaestro/exceptions.py | 4 ++ taskmaestro/yaml_config.py | 12 +++- tests/test_discovery.py | 111 +++++++++++++++++++++++++++++++++++++ 7 files changed, 274 insertions(+), 5 deletions(-) create mode 100644 taskmaestro/discovery.py create mode 100644 tests/test_discovery.py diff --git a/README.md b/README.md index 6c878ac..265cdb7 100644 --- a/README.md +++ b/README.md @@ -196,9 +196,39 @@ grid = GridCase(value=eclipse_case) print(grid.value.name) ``` +## Plugin discovery + +Installed packages can publish tasks and workflows using standard Python entry points: + +```toml +[project.entry-points."taskmaestro.tasks"] +"acme.prepare" = "acme_tasks.prepare:Prepare" + +[project.entry-points."taskmaestro.workflows"] +"acme.analysis" = "acme_tasks.workflows:analysis_workflow" +``` + +A task entry point must resolve to a `Task` subclass and a workflow entry point must +resolve to a `Workflow` instance. Prefix names with the provider name to avoid clashes. +Consumers can discover plugins without scanning package directories: + +```python +from taskmaestro import registered_tasks, registered_workflows + +tasks = registered_tasks() # dict[str, type[Task]] +workflows = registered_workflows() # dict[str, Workflow] +``` + +Use `registered_task_names()` and `registered_workflow_names()` to inspect identifiers +without importing plugin modules, or `get_registered_task(name)` and +`get_registered_workflow(name)` to load one plugin. Duplicate names and invalid plugin +types raise `PluginLoadError`. + ## YAML Configuration -Workflows can be defined entirely in YAML instead of Python. The loader dynamically imports task classes and validates the full configuration: +Workflows can be defined entirely in YAML instead of Python. A `task:` value may be +either a registered task identifier or a dotted Python class path. The loader resolves +registered identifiers first and validates the full configuration: ```yaml # workflow.yaml diff --git a/pyproject.toml b/pyproject.toml index cc46c92..87380f3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "taskmaestro" -version = "0.1.0" +version = "0.2.0" description = "A Python 3.12+ library for defining and executing typed DAG task workflows" readme = "README.md" requires-python = ">=3.12" diff --git a/taskmaestro/__init__.py b/taskmaestro/__init__.py index 017ac8a..f613347 100644 --- a/taskmaestro/__init__.py +++ b/taskmaestro/__init__.py @@ -1,13 +1,24 @@ """Workflow Runner: typed DAG task workflows with Pydantic models.""" -__version__ = "0.1.0" +__version__ = "0.2.0" from taskmaestro.context import ExecutionContext +from taskmaestro.discovery import ( + TASK_ENTRY_POINT_GROUP, + WORKFLOW_ENTRY_POINT_GROUP, + get_registered_task, + get_registered_workflow, + registered_task_names, + registered_tasks, + registered_workflow_names, + registered_workflows, +) from taskmaestro.exceptions import ( ConfigLoadError, CycleDetectedError, IncompleteInputError, JobStateError, + PluginLoadError, TaskExecutionError, TaskOutputTypeError, TaskTimeoutError, @@ -28,6 +39,8 @@ ) __all__ = [ + "TASK_ENTRY_POINT_GROUP", + "WORKFLOW_ENTRY_POINT_GROUP", "ConfigLoadError", "CycleDetectedError", "EmptyConfig", @@ -39,6 +52,7 @@ "JobStatus", "LoadedWorkflow", "ObjectModel", + "PluginLoadError", "Runner", "Task", "TaskExecutionError", @@ -50,7 +64,13 @@ "WorkflowBuilder", "WorkflowDefinitionError", "WorkflowRunnerError", + "get_registered_task", + "get_registered_workflow", "load_workflow_from_yaml", + "registered_task_names", + "registered_tasks", + "registered_workflow_names", + "registered_workflows", "run_workflow_from_yaml", "to_mermaid", "workflow_task", diff --git a/taskmaestro/discovery.py b/taskmaestro/discovery.py new file mode 100644 index 0000000..de5bbe3 --- /dev/null +++ b/taskmaestro/discovery.py @@ -0,0 +1,96 @@ +"""Discovery of tasks and workflows published by installed distributions.""" + +from __future__ import annotations + +from importlib.metadata import EntryPoint, entry_points +from typing import Any + +from taskmaestro.exceptions import PluginLoadError +from taskmaestro.task import Task +from taskmaestro.workflow import Workflow + +TASK_ENTRY_POINT_GROUP = "taskmaestro.tasks" +WORKFLOW_ENTRY_POINT_GROUP = "taskmaestro.workflows" + +def _entry_points_by_name(group: str) -> dict[str, EntryPoint]: + """Return entry points in *group*, rejecting ambiguous registrations.""" + result: dict[str, EntryPoint] = {} + for entry_point in entry_points(group=group): + if entry_point.name in result: + raise PluginLoadError( + f"Multiple entry points named '{entry_point.name}' are registered in '{group}'" + ) + result[entry_point.name] = entry_point + return result + + +def _load[T](entry_point: EntryPoint, expected_type: type[T], kind: str) -> T: + try: + value: Any = entry_point.load() + except Exception as exc: + raise PluginLoadError( + f"Cannot load {kind} entry point '{entry_point.name}' " + f"({entry_point.value}): {exc}" + ) from exc + if not isinstance(value, expected_type): + raise PluginLoadError( + f"{kind.capitalize()} entry point '{entry_point.name}' ({entry_point.value}) " + f"must resolve to a {expected_type.__name__}" + ) + return value + + +def registered_task_names() -> set[str]: + """Return registered task identifiers without importing their packages.""" + return set(_entry_points_by_name(TASK_ENTRY_POINT_GROUP)) + + +def registered_workflow_names() -> set[str]: + """Return registered workflow identifiers without importing their packages.""" + return set(_entry_points_by_name(WORKFLOW_ENTRY_POINT_GROUP)) + + +def registered_tasks() -> dict[str, type[Task[Any, Any]]]: + """Load all tasks registered in the ``taskmaestro.tasks`` entry-point group. + + The returned mapping is keyed by the entry-point name, which is the stable + identifier that configuration files and discovery clients should use. + """ + tasks: dict[str, type[Task[Any, Any]]] = {} + for name, entry_point in _entry_points_by_name(TASK_ENTRY_POINT_GROUP).items(): + task = _load(entry_point, type, "task") + if not issubclass(task, Task): + raise PluginLoadError( + f"Task entry point '{name}' ({entry_point.value}) must resolve to a Task subclass" + ) + tasks[name] = task + return tasks + + +def registered_workflows() -> dict[str, Workflow]: + """Load all workflows registered in the ``taskmaestro.workflows`` group.""" + return { + name: _load(entry_point, Workflow, "workflow") + for name, entry_point in _entry_points_by_name(WORKFLOW_ENTRY_POINT_GROUP).items() + } + + +def get_registered_task(name: str) -> type[Task[Any, Any]]: + """Load one registered task by its entry-point name.""" + entry_point = _entry_points_by_name(TASK_ENTRY_POINT_GROUP).get(name) + if entry_point is None: + raise PluginLoadError(f"No task entry point named '{name}' is registered") + task = _load(entry_point, type, "task") + if not issubclass(task, Task): + raise PluginLoadError( + f"Task entry point '{name}' ({entry_point.value}) must resolve to a Task subclass" + ) + return task + + +def get_registered_workflow(name: str) -> Workflow: + """Load one registered workflow by its entry-point name.""" + entry_point = _entry_points_by_name(WORKFLOW_ENTRY_POINT_GROUP).get(name) + if entry_point is None: + raise PluginLoadError(f"No workflow entry point named '{name}' is registered") + return _load(entry_point, Workflow, "workflow") diff --git a/taskmaestro/exceptions.py b/taskmaestro/exceptions.py index 7e5a701..c5c47b1 100644 --- a/taskmaestro/exceptions.py +++ b/taskmaestro/exceptions.py @@ -35,3 +35,7 @@ class TaskTimeoutError(TaskExecutionError): class ConfigLoadError(WorkflowRunnerError): """Raised when YAML config loading fails (parse errors, import failures, validation).""" + + +class PluginLoadError(WorkflowRunnerError): + """Raised when an installed task or workflow entry point is invalid.""" diff --git a/taskmaestro/yaml_config.py b/taskmaestro/yaml_config.py index c292f8f..3f0b644 100644 --- a/taskmaestro/yaml_config.py +++ b/taskmaestro/yaml_config.py @@ -12,7 +12,8 @@ from pydantic import BaseModel, Field, ValidationError, model_validator from taskmaestro.context import ExecutionContext -from taskmaestro.exceptions import ConfigLoadError +from taskmaestro.discovery import get_registered_task, registered_task_names +from taskmaestro.exceptions import ConfigLoadError, PluginLoadError from taskmaestro.hooks.base import BaseHook from taskmaestro.job import EmptyConfig, Job, JobConfiguration from taskmaestro.runner import Runner @@ -195,6 +196,7 @@ def _load_workflow_only( # 4. Resolve task import paths (handles both task: and workflow: entries) base_dir = workflow_path.parent task_classes: dict[str, type[Task[Any, Any]]] = {} + installed_task_names = registered_task_names() for task_config in config.workflow.tasks: if task_config.workflow: # Recursive workflow reference @@ -209,7 +211,13 @@ def _load_workflow_only( task_classes[task_config.workflow] = wrapped_cls else: assert task_config.task is not None - cls = import_class(task_config.task) + try: + if task_config.task in installed_task_names: + cls = get_registered_task(task_config.task) + else: + cls = import_class(task_config.task) + except PluginLoadError as exc: + raise ConfigLoadError(str(exc)) from exc if not (isinstance(cls, type) and issubclass(cls, Task)): raise ConfigLoadError(f"'{task_config.task}' is not a Task subclass") task_classes[task_config.task] = cls diff --git a/tests/test_discovery.py b/tests/test_discovery.py new file mode 100644 index 0000000..7b8ec52 --- /dev/null +++ b/tests/test_discovery.py @@ -0,0 +1,111 @@ +"""Tests for installed task and workflow discovery.""" + +from __future__ import annotations + +from importlib.metadata import EntryPoint +from pathlib import Path +from typing import Any + +import pytest +from pydantic import BaseModel + +from taskmaestro import ( + ExecutionContext, + PluginLoadError, + Task, + Workflow, + get_registered_task, + get_registered_workflow, + load_workflow_from_yaml, + registered_task_names, + registered_tasks, + registered_workflow_names, + registered_workflows, +) + + +class Input(BaseModel): + value: int + + +class Output(BaseModel): + value: int + + +class ExampleTask(Task[Input, Output]): + def run(self, input: Input, ctx: ExecutionContext) -> Output: + return Output(value=input.value + 1) + + +example_workflow = Workflow("example", [ExampleTask]) + + +def _entry_point(name: str, value: str, group: str) -> EntryPoint: + return EntryPoint(name=name, value=value, group=group) + + +@pytest.fixture +def plugin_entry_points(monkeypatch: pytest.MonkeyPatch) -> None: + entries = { + "taskmaestro.tasks": [ + _entry_point( + "example.increment", "tests.test_discovery:ExampleTask", "taskmaestro.tasks" + ) + ], + "taskmaestro.workflows": [ + _entry_point( + "example.workflow", + "tests.test_discovery:example_workflow", + "taskmaestro.workflows", + ) + ], + } + monkeypatch.setattr( + "taskmaestro.discovery.entry_points", lambda *, group: entries.get(group, []) + ) + + +def test_discovers_registered_plugins(plugin_entry_points: None) -> None: + assert registered_task_names() == {"example.increment"} + assert registered_workflow_names() == {"example.workflow"} + assert registered_tasks() == {"example.increment": ExampleTask} + assert registered_workflows() == {"example.workflow": example_workflow} + assert get_registered_task("example.increment") is ExampleTask + assert get_registered_workflow("example.workflow") is example_workflow + + +def test_registered_task_can_be_used_in_yaml( + plugin_entry_points: None, tmp_path: Path +) -> None: + workflow_path = tmp_path / "workflow.yaml" + workflow_path.write_text( + "workflow:\n" + " name: entry_point_workflow\n" + " tasks:\n" + " - task: example.increment\n" + ) + input_path = tmp_path / "input.yaml" + input_path.write_text("value: 4\n") + + result = load_workflow_from_yaml(workflow_path, input_path).run() + + assert result.result == Output(value=5) + + +def test_rejects_wrong_plugin_type(monkeypatch: pytest.MonkeyPatch) -> None: + entry = _entry_point("invalid", "tests.test_discovery:Input", "taskmaestro.tasks") + monkeypatch.setattr("taskmaestro.discovery.entry_points", lambda *, group: [entry]) + + with pytest.raises(PluginLoadError, match="Task subclass"): + registered_tasks() + + +def test_rejects_duplicate_names(monkeypatch: pytest.MonkeyPatch) -> None: + entries: list[Any] = [ + _entry_point("duplicate", "tests.test_discovery:ExampleTask", "taskmaestro.tasks"), + _entry_point("duplicate", "tests.test_discovery:ExampleTask", "taskmaestro.tasks"), + ] + monkeypatch.setattr("taskmaestro.discovery.entry_points", lambda *, group: entries) + + with pytest.raises(PluginLoadError, match="Multiple entry points"): + registered_tasks() From 65d6d44dfa8f7a2e7369d31be6f34c3f1716855f Mon Sep 17 00:00:00 2001 From: Kristian Bendiksen Date: Wed, 19 Aug 2026 09:01:08 +0200 Subject: [PATCH 2/2] Require Ruff 0.16 and apply formatting --- README.md | 43 ++++++++++++++++++++++++--------- examples/resinsight/pipeline.py | 17 ++++++++++--- pyproject.toml | 2 +- taskmaestro/discovery.py | 4 +-- tests/test_discovery.py | 9 ++----- 5 files changed, 50 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 265cdb7..e439a48 100644 --- a/README.md +++ b/README.md @@ -22,25 +22,30 @@ pip install -e ".[dev]" from pydantic import BaseModel from taskmaestro import Task, Workflow, Job, Runner, ExecutionContext + class NumberInput(BaseModel): value: int + class NumberOutput(BaseModel): value: int + class AddOne(Task[NumberInput, NumberOutput]): def run(self, input: NumberInput, ctx: ExecutionContext) -> NumberOutput: return NumberOutput(value=input.value + 1) + class Double(Task[NumberOutput, NumberOutput]): def run(self, input: NumberOutput, ctx: ExecutionContext) -> NumberOutput: return NumberOutput(value=input.value * 2) + workflow = Workflow(name="math", tasks=[AddOne, Double]) job = Job(workflow=workflow, config=NumberInput(value=5)) result = Runner().run(job) -print(result.status) # "completed" +print(result.status) # "completed" print(result.result.value) # 12 ``` @@ -50,31 +55,39 @@ print(result.result.value) # 12 from pydantic import BaseModel from taskmaestro import Task, Workflow, Job, Runner, ExecutionContext + class Input(BaseModel): value: int + class Output(BaseModel): value: int + class MergedInput(BaseModel): a: Output b: Output + class MergedOutput(BaseModel): total: int + class BranchA(Task[Input, Output]): def run(self, input: Input, ctx: ExecutionContext) -> Output: return Output(value=input.value + 1) + class BranchB(Task[Input, Output]): def run(self, input: Input, ctx: ExecutionContext) -> Output: return Output(value=input.value * 2) + class Merge(Task[MergedInput, MergedOutput]): def run(self, input: MergedInput, ctx: ExecutionContext) -> MergedOutput: return MergedOutput(total=input.a.value + input.b.value) + workflow = ( Workflow.builder(name="fan_in") .add_task(BranchA) @@ -86,7 +99,7 @@ workflow = ( job = Job(workflow=workflow, config=Input(value=5)) result = Runner().run(job) -print(result.status) # "completed" +print(result.status) # "completed" print(result.result.total) # 16 (6 + 10) ``` @@ -131,15 +144,17 @@ from taskmaestro import EmptyConfig, Job, JobConfiguration, Workflow workflow = ( Workflow.builder(name="configured") - .add_task(LoadModel, config_fields=["path"]) # root: all input from config + .add_task(LoadModel, config_fields=["path"]) # root: all input from config .add_task(Transform, depends_on=LoadModel, config_fields=["scale_factor"]) # mixed .build() ) -job_config = JobConfiguration({ - "load_model": {"path": "/data/model.egrid"}, - "transform": {"scale_factor": 2.5}, -}) +job_config = JobConfiguration( + { + "load_model": {"path": "/data/model.egrid"}, + "transform": {"scale_factor": 2.5}, + } +) job = Job(workflow=workflow, config=EmptyConfig(), job_configuration=job_config) result = Runner().run(job) @@ -160,15 +175,19 @@ class ExtractKeywords(Task): def run(self, input: Inputs, ctx: ExecutionContext) -> Outputs: ... + workflow = ( Workflow.builder(name="analysis") .add_task(ExtractKeywords, depends_on=PrepareText) .add_task( BuildReport, depends_on={ - "keywords": (ExtractKeywords, "keywords"), # routes .keywords field - "num_words_removed": (ExtractKeywords, "num_words_removed"), # routes .num_words_removed - "stats": ComputeWordStats, # whole output + "keywords": (ExtractKeywords, "keywords"), # routes .keywords field + "num_words_removed": ( + ExtractKeywords, + "num_words_removed", + ), # routes .num_words_removed + "stats": ComputeWordStats, # whole output }, ) .build() @@ -186,11 +205,13 @@ from taskmaestro import ObjectModel GridCase = ObjectModel[rips.EclipseCase] WellPath = ObjectModel[rips.WellPath] + # Subclass — adds fields alongside the wrapped object class AddPerforationInput(ObjectModel[rips.WellPath]): start_md: float end_md: float + # Access the wrapped object via .value grid = GridCase(value=eclipse_case) print(grid.value.name) @@ -215,7 +236,7 @@ Consumers can discover plugins without scanning package directories: ```python from taskmaestro import registered_tasks, registered_workflows -tasks = registered_tasks() # dict[str, type[Task]] +tasks = registered_tasks() # dict[str, type[Task]] workflows = registered_workflows() # dict[str, Workflow] ``` diff --git a/examples/resinsight/pipeline.py b/examples/resinsight/pipeline.py index 6b773df..dbb5551 100644 --- a/examples/resinsight/pipeline.py +++ b/examples/resinsight/pipeline.py @@ -114,8 +114,12 @@ class AddPerforationInput(BaseModel): resinsight: RipsInstance well_path: WellPath - event_date: str = Field(description="Perforation event date in ISO format (YYYY-MM-DD)", examples=["2024-01-15"]) - start_md: float = Field(description="Start measured depth of the perforation interval (m)", ge=0) + event_date: str = Field( + description="Perforation event date in ISO format (YYYY-MM-DD)", examples=["2024-01-15"] + ) + start_md: float = Field( + description="Start measured depth of the perforation interval (m)", ge=0 + ) end_md: float = Field(description="End measured depth of the perforation interval (m)", ge=0) @@ -221,14 +225,19 @@ class Inputs(BaseModel): grid_case: GridCase perforation_1: PerforationOutput perforation_2: PerforationOutput - event_date: str = Field(description="Timestamp written into the exported schedule file", examples=["2024-05-01"]) + event_date: str = Field( + description="Timestamp written into the exported schedule file", + examples=["2024-05-01"], + ) export_path: str = Field(description="Output path for the .sch completions file") class Outputs(BaseModel): """Exported completions summary.""" export_file: str = Field(description="Path to the generated completions file") - well_path_names: list[str] = Field(description="Names of the well paths included in the export") + well_path_names: list[str] = Field( + description="Names of the well paths included in the export" + ) def run(self, input: Inputs, ctx: ExecutionContext) -> Outputs: eclipse_case = input.grid_case.value diff --git a/pyproject.toml b/pyproject.toml index 87380f3..2e4a407 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ packages = ["taskmaestro"] dev = [ "pytest>=8.0", "pytest-cov>=5.0", - "ruff>=0.4", + "ruff>=0.16", "mypy>=1.10", "types-PyYAML>=6.0", ] diff --git a/taskmaestro/discovery.py b/taskmaestro/discovery.py index de5bbe3..7546686 100644 --- a/taskmaestro/discovery.py +++ b/taskmaestro/discovery.py @@ -12,6 +12,7 @@ TASK_ENTRY_POINT_GROUP = "taskmaestro.tasks" WORKFLOW_ENTRY_POINT_GROUP = "taskmaestro.workflows" + def _entry_points_by_name(group: str) -> dict[str, EntryPoint]: """Return entry points in *group*, rejecting ambiguous registrations.""" result: dict[str, EntryPoint] = {} @@ -29,8 +30,7 @@ def _load[T](entry_point: EntryPoint, expected_type: type[T], kind: str) -> T: value: Any = entry_point.load() except Exception as exc: raise PluginLoadError( - f"Cannot load {kind} entry point '{entry_point.name}' " - f"({entry_point.value}): {exc}" + f"Cannot load {kind} entry point '{entry_point.name}' ({entry_point.value}): {exc}" ) from exc if not isinstance(value, expected_type): raise PluginLoadError( diff --git a/tests/test_discovery.py b/tests/test_discovery.py index 7b8ec52..9f3fe2b 100644 --- a/tests/test_discovery.py +++ b/tests/test_discovery.py @@ -74,15 +74,10 @@ def test_discovers_registered_plugins(plugin_entry_points: None) -> None: assert get_registered_workflow("example.workflow") is example_workflow -def test_registered_task_can_be_used_in_yaml( - plugin_entry_points: None, tmp_path: Path -) -> None: +def test_registered_task_can_be_used_in_yaml(plugin_entry_points: None, tmp_path: Path) -> None: workflow_path = tmp_path / "workflow.yaml" workflow_path.write_text( - "workflow:\n" - " name: entry_point_workflow\n" - " tasks:\n" - " - task: example.increment\n" + "workflow:\n name: entry_point_workflow\n tasks:\n - task: example.increment\n" ) input_path = tmp_path / "input.yaml" input_path.write_text("value: 4\n")