stagegate 0.1.0__py3-none-any.whl

stagegate/__init__.py

@@ -0,0 +1,20 @@
+"""Public package namespace for stage-aware local pipeline execution."""
+
+from .exceptions import CancelledError, UnknownResourceError, UnschedulableTaskError
+from .handles import PipelineHandle, TaskHandle
+from .pipeline import Pipeline
+from .scheduler import Scheduler
+from .wait import ALL_COMPLETED, FIRST_COMPLETED, FIRST_EXCEPTION
+
+__all__ = [
+    "Scheduler",
+    "Pipeline",
+    "TaskHandle",
+    "PipelineHandle",
+    "FIRST_COMPLETED",
+    "FIRST_EXCEPTION",
+    "ALL_COMPLETED",
+    "CancelledError",
+    "UnknownResourceError",
+    "UnschedulableTaskError",
+]

stagegate/_records.py

@@ -0,0 +1,150 @@
+"""Internal scheduler record and queue-entry definitions."""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from ._states import PipelineState, SchedulerState, TaskState
+
+if TYPE_CHECKING:
+    from .pipeline import Pipeline
+    from .scheduler import Scheduler
+
+
+TerminalTaskState = frozenset(
+    {
+        TaskState.SUCCEEDED,
+        TaskState.FAILED,
+        TaskState.CANCELLED,
+    }
+)
+
+TerminalPipelineState = frozenset(
+    {
+        PipelineState.SUCCEEDED,
+        PipelineState.FAILED,
+        PipelineState.CANCELLED,
+    }
+)
+
+
+@dataclass(order=True, frozen=True, slots=True)
+class PipelineQueueEntry:
+    """FIFO queue entry for pipeline admission."""
+
+    enqueue_seq: int
+    record: PipelineRecord = field(compare=False)
+
+
+@dataclass(order=True, frozen=True, slots=True)
+class TaskPriorityKey:
+    """Deterministic ordering key for queued tasks."""
+
+    neg_stage: int
+    pipeline_enqueue_seq: int
+    pipeline_local_task_seq: int
+    global_task_submit_seq: int
+
+
+@dataclass(order=True, frozen=True, slots=True)
+class TaskQueueEntry:
+    """Priority-queue entry for task admission control."""
+
+    priority: TaskPriorityKey
+    record: TaskRecord = field(compare=False)
+
+    @classmethod
+    def from_record(cls, record: TaskRecord) -> TaskQueueEntry:
+        """Build a queue entry from the record's canonical priority key."""
+
+        return cls(priority=record.priority_key(), record=record)
+
+
+@dataclass(order=True, frozen=True, slots=True)
+class ReadyQueueEntry:
+    """FIFO queue entry for tasks admitted but not yet started by a worker."""
+
+    ready_seq: int
+    record: TaskRecord = field(compare=False)
+
+
+@dataclass(slots=True)
+class PipelineRecord:
+    """Mutable source-of-truth record for one pipeline instance."""
+
+    scheduler: Scheduler
+    pipeline: Pipeline
+    pipeline_id: int
+    enqueue_seq: int
+    state: PipelineState = PipelineState.QUEUED
+    stage_index: int = 0
+    next_task_seq: int = 0
+    result_value: Any = None
+    exception: BaseException | None = None
+    coordinator_thread_ident: int | None = None
+    task_records: list[TaskRecord] = field(default_factory=list)
+
+    def is_terminal(self) -> bool:
+        """Return whether the pipeline is in a terminal state."""
+
+        return self.state in TerminalPipelineState
+
+
+@dataclass(slots=True)
+class TaskRecord:
+    """Mutable source-of-truth record for one scheduled task."""
+
+    scheduler: Scheduler
+    pipeline_record: PipelineRecord
+    task_id: int
+    fn: Callable[..., Any]
+    resources_required: dict[str, int | float]
+    args: tuple[Any, ...] = ()
+    kwargs: dict[str, Any] = field(default_factory=dict)
+    name: str | None = None
+    stage_snapshot: int = 0
+    pipeline_enqueue_seq: int = 0
+    pipeline_local_task_seq: int = 0
+    global_task_submit_seq: int = 0
+    state: TaskState = TaskState.QUEUED
+    result_value: Any = None
+    exception: BaseException | None = None
+    ready_seq: int | None = None
+    worker_thread_ident: int | None = None
+
+    def priority_key(self) -> TaskPriorityKey:
+        """Build the ordering key used by the task priority queue."""
+
+        return TaskPriorityKey(
+            neg_stage=-self.stage_snapshot,
+            pipeline_enqueue_seq=self.pipeline_enqueue_seq,
+            pipeline_local_task_seq=self.pipeline_local_task_seq,
+            global_task_submit_seq=self.global_task_submit_seq,
+        )
+
+    def is_terminal(self) -> bool:
+        """Return whether the task is in a terminal state."""
+
+        return self.state in TerminalTaskState
+
+
+@dataclass(slots=True)
+class SchedulerRuntime:
+    """Mutable internal scheduler-wide runtime bookkeeping."""
+
+    state: SchedulerState = SchedulerState.OPEN
+    next_pipeline_id: int = 0
+    next_pipeline_enqueue_seq: int = 0
+    next_task_id: int = 0
+    next_global_task_submit_seq: int = 0
+    next_ready_seq: int = 0
+    admitted_task_count: int = 0
+    pipeline_queue: deque[PipelineQueueEntry] = field(default_factory=deque)
+    task_queue: list[TaskQueueEntry] = field(default_factory=list)
+    ready_queue: deque[ReadyQueueEntry] = field(default_factory=deque)
+    resources_in_use: dict[str, int | float] = field(default_factory=dict)
+    pipeline_records: dict[int, PipelineRecord] = field(default_factory=dict)
+    task_records: dict[int, TaskRecord] = field(default_factory=dict)

stagegate/_states.py

@@ -0,0 +1,34 @@
+"""Internal state definitions for scheduler records."""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class SchedulerState(Enum):
+    """Internal lifecycle for scheduler shutdown progression."""
+
+    OPEN = "open"
+    SHUTTING_DOWN = "shutting_down"
+    CLOSED = "closed"
+
+
+class PipelineState(Enum):
+    """Internal pipeline execution states."""
+
+    QUEUED = "queued"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class TaskState(Enum):
+    """Internal task execution states."""
+
+    QUEUED = "queued"
+    READY = "ready"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"

stagegate/_wait_utils.py

@@ -0,0 +1,97 @@
+"""Internal helpers for owner-scoped wait APIs."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Iterable
+from typing import TypeVar
+
+from ._states import PipelineState, TaskState
+from .handles import PipelineHandle, TaskHandle
+from .wait import ALL_COMPLETED, FIRST_COMPLETED, WAIT_CONDITIONS
+
+HandleT = TypeVar("HandleT", TaskHandle, PipelineHandle)
+
+
+def validate_wait_request(
+    handles: Iterable[HandleT],
+    *,
+    expected_type: type[HandleT],
+    owner_check: Callable[[HandleT], bool],
+    timeout: float | None,
+    return_when: str,
+) -> set[HandleT]:
+    """Normalize and validate public wait() inputs."""
+
+    if timeout is not None and timeout < 0:
+        raise ValueError("timeout must be None or a non-negative number")
+    if return_when not in WAIT_CONDITIONS:
+        raise ValueError("invalid return_when")
+
+    candidates = tuple(handles)
+    if not candidates:
+        raise ValueError("handles must not be empty")
+
+    for handle in candidates:
+        if not isinstance(handle, expected_type):
+            raise TypeError("wrong handle type")
+        if not owner_check(handle):
+            raise ValueError("handle does not belong to this owner")
+
+    normalized = set(candidates)
+    return normalized
+
+
+def split_done_pending(handles: set[HandleT]) -> tuple[set[HandleT], set[HandleT]]:
+    """Partition handles into terminal and non-terminal subsets."""
+
+    done: set[HandleT] = set()
+    pending: set[HandleT] = set()
+    for handle in handles:
+        if handle.done():
+            done.add(handle)
+        else:
+            pending.add(handle)
+    return done, pending
+
+
+def should_return(
+    *,
+    done: set[HandleT],
+    pending: set[HandleT],
+    return_when: str,
+) -> bool:
+    """Evaluate the public wait-return condition."""
+
+    if return_when == ALL_COMPLETED:
+        return not pending
+    if return_when == FIRST_COMPLETED:
+        return bool(done)
+    if any(_is_failed(handle) for handle in done):
+        return True
+    return not pending
+
+
+def monotonic_deadline(timeout: float | None) -> float | None:
+    """Convert a public timeout value into an absolute deadline."""
+
+    if timeout is None:
+        return None
+    return time.monotonic() + timeout
+
+
+def remaining_timeout(deadline: float | None) -> float | None:
+    """Compute remaining timeout for a condition wait."""
+
+    if deadline is None:
+        return None
+    remaining = deadline - time.monotonic()
+    if remaining <= 0:
+        return 0.0
+    return remaining
+
+
+def _is_failed(handle: HandleT) -> bool:
+    if isinstance(handle, TaskHandle):
+        return handle._record.state is TaskState.FAILED
+    return handle._record.state is PipelineState.FAILED

stagegate/exceptions.py

@@ -0,0 +1,13 @@
+"""Public exceptions for stagegate."""
+
+
+class CancelledError(Exception):
+    """Raised when result() or exception() is requested from a cancelled handle."""
+
+
+class UnknownResourceError(ValueError):
+    """Raised when a task requests a resource label unknown to the scheduler."""
+
+
+class UnschedulableTaskError(ValueError):
+    """Raised when a single task can never fit within scheduler capacity."""

stagegate/handles.py

@@ -0,0 +1,216 @@
+"""Public handle skeletons."""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any
+
+from ._records import TerminalPipelineState, TerminalTaskState
+from ._states import PipelineState, TaskState
+from .exceptions import CancelledError
+
+if TYPE_CHECKING:
+    from ._records import PipelineRecord, TaskRecord
+
+
+def _validate_timeout(timeout: float | None) -> float | None:
+    if timeout is None:
+        return None
+    if timeout < 0:
+        raise ValueError("timeout must be None or a non-negative number")
+    return timeout
+
+
+def _remaining_timeout(deadline: float) -> float:
+    remaining = deadline - time.monotonic()
+    if remaining <= 0:
+        raise TimeoutError
+    return remaining
+
+
+class TaskHandle:
+    """Handle for a task submitted from a pipeline."""
+
+    __slots__ = ("_record",)
+
+    def __init__(self, record: TaskRecord) -> None:
+        self._record = record
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TaskHandle):
+            return NotImplemented
+        return (
+            self._record.task_id == other._record.task_id
+            and self._record.scheduler is other._record.scheduler
+        )
+
+    def __hash__(self) -> int:
+        return hash((TaskHandle, id(self._record.scheduler), self._record.task_id))
+
+    def __repr__(self) -> str:
+        return (
+            f"TaskHandle(scheduler=0x{id(self._record.scheduler):x}, "
+            f"task_id={self._record.task_id})"
+        )
+
+    def cancel(self) -> bool:
+        """Cancel the task if it has not started yet."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return scheduler._cancel_task_if_possible_locked(self._record)
+
+    def done(self) -> bool:
+        """Return whether the task is terminal."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state in TerminalTaskState
+
+    def running(self) -> bool:
+        """Return whether the task callable is actively running."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state is TaskState.RUNNING
+
+    def cancelled(self) -> bool:
+        """Return whether the task ended in the cancelled state."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state is TaskState.CANCELLED
+
+    def result(self, timeout: float | None = None) -> Any:
+        """Return the task result or raise its stored terminal outcome."""
+        timeout = _validate_timeout(timeout)
+        scheduler = self._record.scheduler
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with scheduler._condition:
+            while self._record.state not in TerminalTaskState:
+                wait_timeout = (
+                    None if deadline is None else _remaining_timeout(deadline)
+                )
+                scheduler._condition.wait(wait_timeout)
+
+            if self._record.state is TaskState.SUCCEEDED:
+                return self._record.result_value
+            if self._record.state is TaskState.FAILED:
+                assert self._record.exception is not None
+                raise self._record.exception
+            raise CancelledError("result() was requested from a cancelled task handle")
+
+    def exception(self, timeout: float | None = None) -> BaseException | None:
+        """Return the task exception object, if any."""
+        timeout = _validate_timeout(timeout)
+        scheduler = self._record.scheduler
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with scheduler._condition:
+            while self._record.state not in TerminalTaskState:
+                wait_timeout = (
+                    None if deadline is None else _remaining_timeout(deadline)
+                )
+                scheduler._condition.wait(wait_timeout)
+
+            if self._record.state is TaskState.SUCCEEDED:
+                return None
+            if self._record.state is TaskState.FAILED:
+                assert self._record.exception is not None
+                return self._record.exception
+            raise CancelledError(
+                "exception() was requested from a cancelled task handle"
+            )
+
+
+class PipelineHandle:
+    """Handle for a pipeline submitted to a scheduler."""
+
+    __slots__ = ("_record",)
+
+    def __init__(self, record: PipelineRecord) -> None:
+        self._record = record
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, PipelineHandle):
+            return NotImplemented
+        return (
+            self._record.pipeline_id == other._record.pipeline_id
+            and self._record.scheduler is other._record.scheduler
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (PipelineHandle, id(self._record.scheduler), self._record.pipeline_id)
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"PipelineHandle(scheduler=0x{id(self._record.scheduler):x}, "
+            f"pipeline_id={self._record.pipeline_id})"
+        )
+
+    def cancel(self) -> bool:
+        """Cancel the pipeline if it has not started yet."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            if self._record.state is PipelineState.QUEUED:
+                self._record.state = PipelineState.CANCELLED
+                scheduler._notify_state_change_locked()
+                return True
+            return False
+
+    def done(self) -> bool:
+        """Return whether the pipeline is terminal."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state in TerminalPipelineState
+
+    def running(self) -> bool:
+        """Return whether the pipeline ``run()`` method is active."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state is PipelineState.RUNNING
+
+    def cancelled(self) -> bool:
+        """Return whether the pipeline ended in the cancelled state."""
+        scheduler = self._record.scheduler
+        with scheduler._condition:
+            return self._record.state is PipelineState.CANCELLED
+
+    def result(self, timeout: float | None = None) -> Any:
+        """Return the pipeline result or raise its stored terminal outcome."""
+        timeout = _validate_timeout(timeout)
+        scheduler = self._record.scheduler
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with scheduler._condition:
+            while self._record.state not in TerminalPipelineState:
+                wait_timeout = (
+                    None if deadline is None else _remaining_timeout(deadline)
+                )
+                scheduler._condition.wait(wait_timeout)
+
+            if self._record.state is PipelineState.SUCCEEDED:
+                return self._record.result_value
+            if self._record.state is PipelineState.FAILED:
+                assert self._record.exception is not None
+                raise self._record.exception
+            raise CancelledError(
+                "result() was requested from a cancelled pipeline handle"
+            )
+
+    def exception(self, timeout: float | None = None) -> BaseException | None:
+        """Return the pipeline exception object, if any."""
+        timeout = _validate_timeout(timeout)
+        scheduler = self._record.scheduler
+        deadline = None if timeout is None else time.monotonic() + timeout
+        with scheduler._condition:
+            while self._record.state not in TerminalPipelineState:
+                wait_timeout = (
+                    None if deadline is None else _remaining_timeout(deadline)
+                )
+                scheduler._condition.wait(wait_timeout)
+
+            if self._record.state is PipelineState.SUCCEEDED:
+                return None
+            if self._record.state is PipelineState.FAILED:
+                assert self._record.exception is not None
+                return self._record.exception
+            raise CancelledError(
+                "exception() was requested from a cancelled pipeline handle"
+            )

stagegate/pipeline.py

@@ -0,0 +1,122 @@
+"""Pipeline base class and task-builder skeleton."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import threading
+from typing import TYPE_CHECKING, Any
+
+from .handles import TaskHandle
+from ._states import PipelineState
+from ._wait_utils import (
+    monotonic_deadline,
+    remaining_timeout,
+    should_return,
+    split_done_pending,
+    validate_wait_request,
+)
+from .wait import ALL_COMPLETED
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable
+
+    from ._records import PipelineRecord
+    from .scheduler import Scheduler
+
+
+@dataclass(frozen=True, slots=True)
+class TaskBuilder:
+    """Factory object returned by ``Pipeline.task(...)``."""
+
+    pipeline: Pipeline
+    fn: Callable[..., Any]
+    resources: dict[str, int | float]
+    args: tuple[Any, ...] = ()
+    kwargs: dict[str, Any] = field(default_factory=dict)
+    name: str | None = None
+
+    def run(self) -> TaskHandle:
+        """Submit the task to the owning scheduler and return its handle."""
+        pipeline = self.pipeline
+        scheduler, _ = pipeline._require_control_context()
+        with scheduler._condition:
+            return scheduler._submit_task_builder_locked(self)
+
+
+class Pipeline:
+    """Base class for user-defined pipelines."""
+
+    _stagegate_record: PipelineRecord | None = None
+    _stagegate_scheduler: Scheduler | None = None
+    _stagegate_submitted: bool = False
+
+    def run(self) -> Any:
+        """Execute pipeline logic on a scheduler-owned coordinator thread."""
+        raise NotImplementedError
+
+    def _require_control_context(self):
+        record = getattr(self, "_stagegate_record", None)
+        scheduler = getattr(self, "_stagegate_scheduler", None)
+        if record is None or scheduler is None:
+            raise RuntimeError("pipeline control requires a running pipeline")
+        if record.state is not PipelineState.RUNNING:
+            raise RuntimeError("pipeline control requires a running pipeline")
+        if record.coordinator_thread_ident != threading.get_ident():
+            raise RuntimeError(
+                "pipeline control is allowed only on the coordinator thread"
+            )
+        return scheduler, record
+
+    def task(
+        self,
+        fn: Callable[..., Any],
+        *,
+        resources: dict[str, int | float],
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+        name: str | None = None,
+    ) -> TaskBuilder:
+        """Create a task builder for later submission via ``.run()``."""
+        return TaskBuilder(
+            pipeline=self,
+            fn=fn,
+            resources=dict(resources),
+            args=args,
+            kwargs={} if kwargs is None else dict(kwargs),
+            name=name,
+        )
+
+    def stage_forward(self) -> None:
+        """Advance the pipeline stage used by future task submissions."""
+        scheduler, record = self._require_control_context()
+        with scheduler._condition:
+            record.stage_index += 1
+
+    def wait(
+        self,
+        handles: Iterable[TaskHandle],
+        timeout: float | None = None,
+        return_when: str = ALL_COMPLETED,
+    ) -> tuple[set[TaskHandle], set[TaskHandle]]:
+        """Wait for task handles created by this pipeline."""
+        # Concrete implementation must validate return_when against WAIT_CONDITIONS.
+        scheduler, _ = self._require_control_context()
+        normalized = validate_wait_request(
+            handles,
+            expected_type=TaskHandle,
+            owner_check=lambda handle: handle._record.pipeline_record.pipeline is self,
+            timeout=timeout,
+            return_when=return_when,
+        )
+        deadline = monotonic_deadline(timeout)
+
+        with scheduler._condition:
+            while True:
+                done, pending = split_done_pending(normalized)
+                if should_return(done=done, pending=pending, return_when=return_when):
+                    return done, pending
+
+                wait_timeout = remaining_timeout(deadline)
+                if wait_timeout == 0.0:
+                    return done, pending
+                scheduler._condition.wait(wait_timeout)