taskmaestro 0.1.0__py3-none-any.whl

taskmaestro/__init__.py

@@ -0,0 +1,57 @@
+"""Workflow Runner: typed DAG task workflows with Pydantic models."""
+
+__version__ = "0.1.0"
+
+from taskmaestro.context import ExecutionContext
+from taskmaestro.exceptions import (
+    ConfigLoadError,
+    CycleDetectedError,
+    IncompleteInputError,
+    JobStateError,
+    TaskExecutionError,
+    TaskOutputTypeError,
+    TaskTimeoutError,
+    WorkflowDefinitionError,
+    WorkflowRunnerError,
+)
+from taskmaestro.job import EmptyConfig, Job, JobConfiguration, JobStatus, TaskResult, TaskStatus
+from taskmaestro.object_model import ObjectModel
+from taskmaestro.runner import Runner
+from taskmaestro.task import Task
+from taskmaestro.visualization import to_mermaid
+from taskmaestro.workflow import Workflow, WorkflowBuilder
+from taskmaestro.workflow_task import workflow_task
+from taskmaestro.yaml_config import (
+    LoadedWorkflow,
+    load_workflow_from_yaml,
+    run_workflow_from_yaml,
+)
+
+__all__ = [
+    "ConfigLoadError",
+    "CycleDetectedError",
+    "EmptyConfig",
+    "ExecutionContext",
+    "IncompleteInputError",
+    "Job",
+    "JobConfiguration",
+    "JobStateError",
+    "JobStatus",
+    "LoadedWorkflow",
+    "ObjectModel",
+    "Runner",
+    "Task",
+    "TaskExecutionError",
+    "TaskOutputTypeError",
+    "TaskResult",
+    "TaskStatus",
+    "TaskTimeoutError",
+    "Workflow",
+    "WorkflowBuilder",
+    "WorkflowDefinitionError",
+    "WorkflowRunnerError",
+    "load_workflow_from_yaml",
+    "run_workflow_from_yaml",
+    "to_mermaid",
+    "workflow_task",
+]

taskmaestro/context.py

@@ -0,0 +1,37 @@
+"""Execution context passed to every task during execution."""
+
+from __future__ import annotations
+
+import logging
+import tempfile
+import uuid
+from pathlib import Path
+from typing import Any
+
+
+class ExecutionContext:
+    """Cross-cutting state passed to every task during execution.
+
+    Carries a correlation ID, logger, scratch directory, and a simple
+    service registry for injecting dependencies (DB connections, HTTP
+    clients, etc.) into tasks.
+    """
+
+    def __init__(
+        self,
+        correlation_id: str | None = None,
+        logger: logging.Logger | None = None,
+        scratch_dir: Path | None = None,
+    ) -> None:
+        self.correlation_id = correlation_id or str(uuid.uuid4())
+        self.logger = logger or logging.getLogger("taskmaestro")
+        self.scratch_dir = scratch_dir or Path(tempfile.gettempdir()) / self.correlation_id
+        self._registry: dict[str, Any] = {}
+
+    def register(self, key: str, service: Any) -> None:
+        """Register a service by key."""
+        self._registry[key] = service
+
+    def resolve(self, key: str) -> Any:
+        """Retrieve a registered service. Raises KeyError if not found."""
+        return self._registry[key]

taskmaestro/exceptions.py

@@ -0,0 +1,37 @@
+"""Exception hierarchy for the workflow runner library."""
+
+
+class WorkflowRunnerError(Exception):
+    """Base exception for all workflow runner errors."""
+
+
+class WorkflowDefinitionError(WorkflowRunnerError):
+    """Raised at workflow construction time for invalid definitions."""
+
+
+class CycleDetectedError(WorkflowDefinitionError):
+    """Dependency graph contains a cycle."""
+
+
+class IncompleteInputError(WorkflowDefinitionError):
+    """Fan-in input has required fields not mapped to any upstream task."""
+
+
+class JobStateError(WorkflowRunnerError):
+    """Job not in expected state (e.g., attempting to re-run a completed job)."""
+
+
+class TaskExecutionError(WorkflowRunnerError):
+    """Raised during task execution."""
+
+
+class TaskOutputTypeError(TaskExecutionError):
+    """Task returned an output whose type doesn't match the declared output type."""
+
+
+class TaskTimeoutError(TaskExecutionError):
+    """Raised when a task exceeds its timeout_seconds."""
+
+
+class ConfigLoadError(WorkflowRunnerError):
+    """Raised when YAML config loading fails (parse errors, import failures, validation)."""

taskmaestro/hooks/__init__.py

@@ -0,0 +1,15 @@
+"""Built-in hooks for the workflow runner."""
+
+from taskmaestro.hooks.base import BaseHook, Event, Hook
+from taskmaestro.hooks.logging import LoggingHook
+from taskmaestro.hooks.persistence import ResultPersistenceHook
+from taskmaestro.hooks.timing import TimingHook
+
+__all__ = [
+    "BaseHook",
+    "Event",
+    "Hook",
+    "LoggingHook",
+    "ResultPersistenceHook",
+    "TimingHook",
+]

taskmaestro/hooks/base.py

@@ -0,0 +1,57 @@
+"""Hook protocol and base implementation for lifecycle events."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+    from taskmaestro.job import Job
+    from taskmaestro.task import Task
+
+
+class Event(StrEnum):
+    """Lifecycle events emitted by the runner."""
+
+    JOB_START = "job_start"
+    JOB_COMPLETE = "job_complete"
+    JOB_FAIL = "job_fail"
+    TASK_START = "task_start"
+    TASK_COMPLETE = "task_complete"
+    TASK_FAIL = "task_fail"
+
+
+@runtime_checkable
+class Hook(Protocol):
+    """Protocol for lifecycle event hooks."""
+
+    def on_job_start(self, job: Job[Any]) -> None: ...
+    def on_job_complete(self, job: Job[Any]) -> None: ...
+    def on_job_fail(self, job: Job[Any]) -> None: ...
+    def on_task_start(self, job: Job[Any], task: Task[Any, Any]) -> None: ...
+    def on_task_complete(self, job: Job[Any], task: Task[Any, Any], output: BaseModel) -> None: ...
+    def on_task_fail(self, job: Job[Any], task: Task[Any, Any], error: Exception) -> None: ...
+
+
+class BaseHook:
+    """Base hook with no-op defaults for all events."""
+
+    def on_job_start(self, job: Job[Any]) -> None:
+        pass
+
+    def on_job_complete(self, job: Job[Any]) -> None:
+        pass
+
+    def on_job_fail(self, job: Job[Any]) -> None:
+        pass
+
+    def on_task_start(self, job: Job[Any], task: Task[Any, Any]) -> None:
+        pass
+
+    def on_task_complete(self, job: Job[Any], task: Task[Any, Any], output: BaseModel) -> None:
+        pass
+
+    def on_task_fail(self, job: Job[Any], task: Task[Any, Any], error: Exception) -> None:
+        pass

taskmaestro/hooks/logging.py

@@ -0,0 +1,44 @@
+"""Logging hook that logs lifecycle events via Python logging."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from pydantic import BaseModel
+
+from taskmaestro.hooks.base import BaseHook
+from taskmaestro.job import Job
+from taskmaestro.task import Task
+
+
+class LoggingHook(BaseHook):
+    """Logs all lifecycle events via Python's logging module."""
+
+    def __init__(self, level: int = logging.INFO) -> None:
+        self._level = level
+        self._logger = logging.getLogger("taskmaestro.hooks.logging")
+
+    def on_job_start(self, job: Job[Any]) -> None:
+        self._logger.log(self._level, "Job started: workflow=%s", job.workflow.name)
+
+    def on_job_complete(self, job: Job[Any]) -> None:
+        self._logger.log(self._level, "Job completed: workflow=%s", job.workflow.name)
+
+    def on_job_fail(self, job: Job[Any]) -> None:
+        self._logger.log(
+            self._level,
+            "Job failed: workflow=%s, failed_task=%s, error=%s",
+            job.workflow.name,
+            job.failed_task,
+            job.error,
+        )
+
+    def on_task_start(self, job: Job[Any], task: Task[Any, Any]) -> None:
+        self._logger.log(self._level, "Task started: %s", task.name)
+
+    def on_task_complete(self, job: Job[Any], task: Task[Any, Any], output: BaseModel) -> None:
+        self._logger.log(self._level, "Task completed: %s", task.name)
+
+    def on_task_fail(self, job: Job[Any], task: Task[Any, Any], error: Exception) -> None:
+        self._logger.log(self._level, "Task failed: %s, error=%s", task.name, error)

taskmaestro/hooks/persistence.py

@@ -0,0 +1,24 @@
+"""Result persistence hook that writes task outputs to JSON files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+from taskmaestro.hooks.base import BaseHook
+from taskmaestro.job import Job
+from taskmaestro.task import Task
+
+
+class ResultPersistenceHook(BaseHook):
+    """Writes {task_name}.json per completed task to an output directory."""
+
+    def __init__(self, output_dir: Path) -> None:
+        self.output_dir = output_dir
+
+    def on_task_complete(self, job: Job[Any], task: Task[Any, Any], output: BaseModel) -> None:
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        output_path = self.output_dir / f"{task.name}.json"
+        output_path.write_text(output.model_dump_json(indent=2))

taskmaestro/hooks/timing.py

@@ -0,0 +1,46 @@
+"""Timing hook that records wall-clock durations for jobs and tasks."""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+from pydantic import BaseModel
+
+from taskmaestro.hooks.base import BaseHook
+from taskmaestro.job import Job
+from taskmaestro.task import Task
+
+
+class TimingHook(BaseHook):
+    """Records wall-clock duration per task and total job time via time.monotonic()."""
+
+    def __init__(self) -> None:
+        self.job_duration: float | None = None
+        self.task_timings: dict[str, float] = {}
+        self._job_start: float | None = None
+        self._task_starts: dict[str, float] = {}
+
+    def on_job_start(self, job: Job[Any]) -> None:
+        self._job_start = time.monotonic()
+
+    def on_job_complete(self, job: Job[Any]) -> None:
+        if self._job_start is not None:
+            self.job_duration = time.monotonic() - self._job_start
+
+    def on_job_fail(self, job: Job[Any]) -> None:
+        if self._job_start is not None:
+            self.job_duration = time.monotonic() - self._job_start
+
+    def on_task_start(self, job: Job[Any], task: Task[Any, Any]) -> None:
+        self._task_starts[task.name] = time.monotonic()
+
+    def on_task_complete(self, job: Job[Any], task: Task[Any, Any], output: BaseModel) -> None:
+        start = self._task_starts.get(task.name)
+        if start is not None:
+            self.task_timings[task.name] = time.monotonic() - start
+
+    def on_task_fail(self, job: Job[Any], task: Task[Any, Any], error: Exception) -> None:
+        start = self._task_starts.get(task.name)
+        if start is not None:
+            self.task_timings[task.name] = time.monotonic() - start

taskmaestro/job.py

@@ -0,0 +1,116 @@
+"""Job: a workflow bound to a specific config, ready to execute."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from enum import StrEnum
+from typing import Any, Generic, TypeVar
+
+from pydantic import BaseModel
+
+from taskmaestro.exceptions import WorkflowDefinitionError
+from taskmaestro.task import get_input_type
+from taskmaestro.workflow import Workflow
+
+
+class JobStatus(StrEnum):
+    """Status of a job."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class TaskStatus(StrEnum):
+    """Status of an individual task execution."""
+
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+@dataclass
+class TaskResult:
+    """Record of a single task's execution within a job."""
+
+    task_name: str
+    status: TaskStatus
+    output: BaseModel | None
+    started_at: datetime
+    duration_seconds: float
+    error: str | None = None
+
+
+class EmptyConfig(BaseModel):
+    """Sentinel config for workflows where all root tasks use JobConfiguration."""
+
+
+class JobConfiguration:
+    """Per-task configuration values, mapping task names to config field dicts.
+
+    Used to provide static configuration values (from YAML or code) that get
+    merged with upstream outputs when constructing task inputs.
+    """
+
+    def __init__(self, config: dict[str, dict[str, Any]]) -> None:
+        self._config = config
+
+    def get_config_for_task(self, name: str) -> dict[str, Any]:
+        """Return config values for a task, or empty dict if none."""
+        return dict(self._config.get(name, {}))
+
+    def configured_tasks(self) -> set[str]:
+        """Return set of task names that have configuration."""
+        return set(self._config.keys())
+
+    def config_fields_for_task(self, name: str) -> set[str]:
+        """Return set of field names configured for a task."""
+        return set(self._config.get(name, {}).keys())
+
+
+C = TypeVar("C", bound=BaseModel)
+
+
+class Job(Generic[C]):
+    """A workflow bound to a specific config, ready to execute.
+
+    Generic over C so that job.config retains its concrete type.
+    Validates that the config type matches the input type of all root tasks.
+    """
+
+    def __init__(
+        self,
+        workflow: Workflow,
+        config: C,
+        *,
+        job_configuration: JobConfiguration | None = None,
+    ) -> None:
+        self.workflow = workflow
+        self.config: C = config
+        self.job_configuration = job_configuration
+        self.status: JobStatus = JobStatus.PENDING
+        self.result: BaseModel | None = None
+        self.error: str | None = None
+        self.failed_task: str | None = None
+        self.started_at: datetime | None = None
+        self.completed_at: datetime | None = None
+        self.task_results: list[TaskResult] = []
+
+        self._validate_root_task_inputs(config)
+
+    def _validate_root_task_inputs(self, config: C) -> None:
+        """Validate that config type matches the input type of all root tasks."""
+        for task_name, deps in self.workflow._dependencies.items():
+            if deps is None:
+                # Skip validation for root tasks that have config_fields
+                config_fields = self.workflow.get_config_fields(task_name)
+                if config_fields:
+                    continue
+                task_cls = self.workflow._tasks[task_name]
+                expected_input = get_input_type(task_cls)
+                if not isinstance(config, expected_input):
+                    raise WorkflowDefinitionError(
+                        f"Root task '{task_name}' expects input type "
+                        f"{expected_input.__name__} but got {type(config).__name__}"
+                    )

taskmaestro/object_model.py

@@ -0,0 +1,19 @@
+"""Generic base model for wrapping arbitrary (non-Pydantic) objects."""
+
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+T = TypeVar("T")
+
+
+class ObjectModel(BaseModel, Generic[T]):
+    """Base model for wrapping arbitrary (non-Pydantic) objects.
+
+    Enables arbitrary_types_allowed so fields can hold native library
+    objects like database connections, API clients, etc.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    value: T

taskmaestro/runner.py

@@ -0,0 +1,201 @@
+"""Runner: the synchronous execution engine for workflows."""
+
+from __future__ import annotations
+
+import signal
+import warnings
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel
+
+from taskmaestro.context import ExecutionContext
+from taskmaestro.exceptions import (
+    JobStateError,
+    TaskOutputTypeError,
+    TaskTimeoutError,
+)
+from taskmaestro.hooks.base import BaseHook, Event
+from taskmaestro.job import Job, JobStatus, TaskResult, TaskStatus
+from taskmaestro.task import get_input_type, get_output_type
+
+
+class Runner:
+    """Synchronous execution engine for workflows.
+
+    Iterates tasks in topological order, assembling each task's input
+    from the outputs of its upstream dependencies.
+    """
+
+    def __init__(self, hooks: list[BaseHook] | None = None) -> None:
+        self.hooks: list[BaseHook] = hooks or []
+
+    def run(
+        self,
+        job: Job[Any],
+        ctx: ExecutionContext | None = None,
+        timeout_seconds: float | None = None,
+    ) -> Job[Any]:
+        """Execute all tasks in topological order. Stops on first failure."""
+        if job.status != JobStatus.PENDING:
+            raise JobStateError(f"Cannot run job with status '{job.status}'; expected 'pending'")
+
+        ctx = ctx or ExecutionContext()
+        workflow = job.workflow
+
+        self._emit(Event.JOB_START, job)
+        job.status = JobStatus.RUNNING
+        job.started_at = datetime.now()
+
+        # Set up job-level timeout
+        job_alarm_set = False
+        if timeout_seconds is not None:
+            job_alarm_set = self._set_alarm(timeout_seconds, "Job")
+
+        outputs: dict[str, BaseModel] = {}
+        job_config = job.job_configuration
+
+        try:
+            for task_name, task_cls in workflow.topological_order():
+                task = task_cls()
+                task.name = task_name  # instance-level override for named instances
+                deps = workflow.get_dependencies(task_name)
+                config_fields = workflow.get_config_fields(task_name)
+                config_values = (
+                    job_config.get_config_for_task(task_name)
+                    if job_config and config_fields
+                    else {}
+                )
+
+                # Assemble input based on dependency type
+                if deps is None:
+                    if config_values:
+                        # Root task with config: build input from config values
+                        input_type = get_input_type(task_cls)
+                        task_input = input_type.model_validate(config_values)
+                    else:
+                        task_input = job.config
+                elif isinstance(deps, str):
+                    if config_values:
+                        # Single dep with config: decompose upstream, merge with config
+                        input_type = get_input_type(task_cls)
+                        upstream_data = outputs[deps].model_dump()
+                        down_fields = input_type.model_fields
+                        merged: dict[str, object] = {
+                            k: v for k, v in upstream_data.items() if k in down_fields
+                        }
+                        merged.update(config_values)
+                        task_input = input_type.model_validate(merged)
+                    else:
+                        task_input = outputs[deps]
+                elif isinstance(deps, tuple):
+                    upstream_name, field_name = deps
+                    task_input = getattr(outputs[upstream_name], field_name)
+                elif isinstance(deps, dict):
+                    input_type = get_input_type(task_cls)
+                    field_values: dict[str, object] = {}
+                    for fname, upstream_ref in deps.items():
+                        if isinstance(upstream_ref, tuple):
+                            up_name, up_field = upstream_ref
+                            field_values[fname] = getattr(outputs[up_name], up_field)
+                        else:
+                            field_values[fname] = outputs[upstream_ref]
+                    if config_values:
+                        field_values.update(config_values)
+                    task_input = input_type.model_validate(field_values)
+                else:
+                    task_input = job.config  # pragma: no cover
+
+                task_started = datetime.now()
+                self._emit(Event.TASK_START, job, task)
+
+                # Set up per-task timeout
+                task_alarm_set = False
+                if task.timeout_seconds is not None:
+                    task_alarm_set = self._set_alarm(task.timeout_seconds, task.name)
+
+                try:
+                    output = task.run(task_input, ctx)
+
+                    # Validate output matches declared type
+                    expected_output_type = get_output_type(task_cls)
+                    if not isinstance(output, expected_output_type):
+                        raise TaskOutputTypeError(
+                            f"Task '{task.name}' returned {type(output).__name__}, "
+                            f"expected {expected_output_type.__name__}"
+                        )
+
+                    duration = (datetime.now() - task_started).total_seconds()
+                    outputs[task.name] = output
+                    job.task_results.append(
+                        TaskResult(
+                            task_name=task.name,
+                            status=TaskStatus.COMPLETED,
+                            output=output,
+                            started_at=task_started,
+                            duration_seconds=duration,
+                        )
+                    )
+                    self._emit(Event.TASK_COMPLETE, job, task, output)
+                except Exception as exc:
+                    duration = (datetime.now() - task_started).total_seconds()
+                    job.status = JobStatus.FAILED
+                    job.error = str(exc)
+                    job.failed_task = task.name
+                    job.completed_at = datetime.now()
+                    job.task_results.append(
+                        TaskResult(
+                            task_name=task.name,
+                            status=TaskStatus.FAILED,
+                            output=None,
+                            started_at=task_started,
+                            duration_seconds=duration,
+                            error=str(exc),
+                        )
+                    )
+                    self._emit(Event.TASK_FAIL, job, task, exc)
+                    self._emit(Event.JOB_FAIL, job)
+                    return job
+                finally:
+                    if task_alarm_set:
+                        signal.alarm(0)
+        finally:
+            if job_alarm_set:
+                signal.alarm(0)
+
+        job.status = JobStatus.COMPLETED
+        job.result = outputs[workflow.result_task_name]
+        job.completed_at = datetime.now()
+        self._emit(Event.JOB_COMPLETE, job)
+        return job
+
+    def _set_alarm(self, seconds: float, label: str) -> bool:
+        """Set a signal.alarm for timeout. Returns True if alarm was set."""
+        try:
+
+            def _handler(signum: int, frame: Any) -> None:
+                raise TaskTimeoutError(f"{label} timed out after {seconds}s")
+
+            signal.signal(signal.SIGALRM, _handler)
+            signal.alarm(int(seconds) if seconds >= 1 else 1)
+            return True
+        except (AttributeError, OSError):
+            warnings.warn(
+                f"signal.alarm not available on this platform; "
+                f"timeout for {label} will not be enforced",
+                stacklevel=2,
+            )
+            return False
+
+    def _emit(self, event: Event, *args: object) -> None:
+        """Dispatch event to all hooks, swallowing any hook errors."""
+        for hook in self.hooks:
+            handler = getattr(hook, f"on_{event}", None)
+            if handler is not None:
+                try:
+                    handler(*args)
+                except Exception:
+                    warnings.warn(
+                        f"Hook {type(hook).__name__} raised during {event}",
+                        stacklevel=2,
+                    )