Add distributed workload protocol

scimesh/distributed/__init__.py

@@ -0,0 +1,28 @@
+"""Coordinator-independent contracts for distributed SciMesh workloads.
+
+This package defines the typed plan and reduction boundary shared by future
+planners, worker adapters, and coordinator bridges.  It intentionally has no
+network, database, or coordinator imports.
+"""
+
+from .models import (
+    ArtifactReference,
+    CompletedPartial,
+    DistributedPlan,
+    FinalResult,
+    PlannedTask,
+)
+from .registry import DistributedWorkloadRegistry, PlanningService, WorkloadDescription
+from .workload import DistributedWorkload
+
+__all__ = [
+    "ArtifactReference",
+    "CompletedPartial",
+    "DistributedPlan",
+    "DistributedWorkload",
+    "DistributedWorkloadRegistry",
+    "FinalResult",
+    "PlannedTask",
+    "PlanningService",
+    "WorkloadDescription",
+]

scimesh/distributed/models.py

@@ -0,0 +1,262 @@
+"""Versioned, JSON-safe value objects for distributed workload contracts."""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+from dataclasses import dataclass
+from typing import Any, Mapping, Sequence
+from uuid import UUID
+
+
+SCHEMA_VERSION = 1
+_WORKLOAD_NAME = re.compile(r"^[a-z][a-z0-9]*(?:-[a-z0-9]+)*$")
+
+
+def _canonical_uuid(value: object, field: str) -> str:
+    if not isinstance(value, str):
+        raise ValueError(f"{field} must be a UUID string")
+    try:
+        return str(UUID(value))
+    except ValueError as error:
+        raise ValueError(f"{field} must be a UUID string") from error
+
+
+def _sha256(value: object, field: str) -> str:
+    if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value):
+        raise ValueError(f"{field} must be a lowercase SHA-256 hex digest")
+    return value
+
+
+def _content_type(value: object, field: str) -> str:
+    if not isinstance(value, str) or not value or len(value) > 128:
+        raise ValueError(f"{field} must be a non-empty content type")
+    if any(character.isspace() or ord(character) < 32 for character in value):
+        raise ValueError(f"{field} must be a non-empty content type")
+    return value
+
+
+def _workload_name(value: object, field: str = "workload") -> str:
+    if not isinstance(value, str) or not _WORKLOAD_NAME.fullmatch(value):
+        raise ValueError(f"{field} must be a canonical hyphenated workload name")
+    return value
+
+
+def _json_value(value: object, field: str) -> Any:
+    """Deep-copy a JSON value and reject non-finite or non-string-key data."""
+    if value is None or isinstance(value, (bool, int)):
+        return value
+    if isinstance(value, str):
+        # Coordinator-owned artifacts are represented exclusively by
+        # ArtifactReference. A URI or a local path in a generic JSON payload
+        # would let a planner accidentally leak a bridge/worker implementation
+        # detail into durable task metadata.
+        forbidden_prefixes = ("file://", "worker://", "http://", "https://", "s3://", "/")
+        is_windows_path = len(value) >= 3 and value[0].isalpha() and value[1:3] in (":/", ":\\")
+        if value.startswith(forbidden_prefixes) or is_windows_path:
+            raise ValueError(f"{field} must not contain a URI or local path")
+        return value
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            raise ValueError(f"{field} must not contain NaN or infinity")
+        return value
+    if isinstance(value, Mapping):
+        copied: dict[str, Any] = {}
+        for key, child in value.items():
+            if not isinstance(key, str):
+                raise ValueError(f"{field} must use string object keys")
+            copied[key] = _json_value(child, f"{field}.{key}")
+        return copied
+    if isinstance(value, (list, tuple)):
+        return [_json_value(child, f"{field}[]") for child in value]
+    raise ValueError(f"{field} must contain only JSON-compatible values")
+
+
+def _json_mapping(value: object, field: str) -> dict[str, Any]:
+    if not isinstance(value, Mapping):
+        raise ValueError(f"{field} must be an object")
+    return _json_value(value, field)
+
+
+@dataclass(frozen=True)
+class ArtifactReference:
+    """Immutable coordinator-owned artifact identity used in a plan."""
+
+    artifact_id: str
+    sha256: str
+    content_type: str
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "artifact_id", _canonical_uuid(self.artifact_id, "artifact_id"))
+        object.__setattr__(self, "sha256", _sha256(self.sha256, "sha256"))
+        object.__setattr__(self, "content_type", _content_type(self.content_type, "content_type"))
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "artifact_id": self.artifact_id,
+            "sha256": self.sha256,
+            "content_type": self.content_type,
+        }
+
+    @classmethod
+    def from_dict(cls, value: object) -> "ArtifactReference":
+        if not isinstance(value, Mapping):
+            raise ValueError("artifact reference must be an object")
+        _require_exact_keys(value, {"artifact_id", "sha256", "content_type"}, "artifact reference")
+        return cls(
+            artifact_id=value["artifact_id"],
+            sha256=value["sha256"],
+            content_type=value["content_type"],
+        )
+
+
+@dataclass(frozen=True)
+class PlannedTask:
+    """One deterministically indexed, artifact-backed worker task."""
+
+    chunk_index: int
+    input_artifact: ArtifactReference
+    parameters: Mapping[str, object]
+
+    def __post_init__(self) -> None:
+        if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
+            raise ValueError("chunk_index must be a non-negative integer")
+        if not isinstance(self.input_artifact, ArtifactReference):
+            raise ValueError("input_artifact must be an ArtifactReference")
+        object.__setattr__(self, "parameters", _json_mapping(self.parameters, "task parameters"))
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "chunk_index": self.chunk_index,
+            "input_artifact": self.input_artifact.to_dict(),
+            "parameters": _json_value(self.parameters, "task parameters"),
+        }
+
+    @classmethod
+    def from_dict(cls, value: object) -> "PlannedTask":
+        if not isinstance(value, Mapping):
+            raise ValueError("planned task must be an object")
+        _require_exact_keys(value, {"chunk_index", "input_artifact", "parameters"}, "planned task")
+        return cls(
+            chunk_index=value["chunk_index"],
+            input_artifact=ArtifactReference.from_dict(value["input_artifact"]),
+            parameters=value["parameters"],
+        )
+
+
+@dataclass(frozen=True)
+class DistributedPlan:
+    """The complete schema-versioned output of a distributed planner."""
+
+    workload: str
+    resolved_parameters: Mapping[str, object]
+    tasks: Sequence[PlannedTask]
+    schema_version: int = SCHEMA_VERSION
+
+    def __post_init__(self) -> None:
+        if self.schema_version != SCHEMA_VERSION:
+            raise ValueError(f"schema_version must be {SCHEMA_VERSION}")
+        object.__setattr__(self, "workload", _workload_name(self.workload))
+        object.__setattr__(self, "resolved_parameters", _json_mapping(self.resolved_parameters, "resolved_parameters"))
+        task_list = tuple(self.tasks)
+        if not task_list:
+            raise ValueError("plan must contain at least one task")
+        if any(not isinstance(task, PlannedTask) for task in task_list):
+            raise ValueError("tasks must contain PlannedTask values")
+        indexes = [task.chunk_index for task in task_list]
+        if indexes != sorted(indexes) or len(set(indexes)) != len(indexes):
+            raise ValueError("tasks must have unique, ascending chunk_index values")
+        object.__setattr__(self, "tasks", task_list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "workload": self.workload,
+            "resolved_parameters": _json_value(self.resolved_parameters, "resolved_parameters"),
+            "tasks": [task.to_dict() for task in self.tasks],
+        }
+
+    def to_json(self) -> str:
+        """Return stable JSON suitable for hashing, tests, and durable payloads."""
+        return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False)
+
+    @classmethod
+    def from_dict(cls, value: object) -> "DistributedPlan":
+        if not isinstance(value, Mapping):
+            raise ValueError("distributed plan must be an object")
+        _require_exact_keys(
+            value,
+            {"schema_version", "workload", "resolved_parameters", "tasks"},
+            "distributed plan",
+        )
+        raw_tasks = value["tasks"]
+        if not isinstance(raw_tasks, list):
+            raise ValueError("tasks must be an array")
+        return cls(
+            schema_version=value["schema_version"],
+            workload=value["workload"],
+            resolved_parameters=value["resolved_parameters"],
+            tasks=tuple(PlannedTask.from_dict(task) for task in raw_tasks),
+        )
+
+    @classmethod
+    def from_json(cls, value: str) -> "DistributedPlan":
+        try:
+            decoded = json.loads(value)
+        except (TypeError, json.JSONDecodeError) as error:
+            raise ValueError("distributed plan must be valid JSON") from error
+        return cls.from_dict(decoded)
+
+
+@dataclass(frozen=True)
+class CompletedPartial:
+    """Coordinator-owned partial output supplied to a reducer."""
+
+    chunk_index: int
+    artifact: ArtifactReference
+    metrics: Mapping[str, int | float]
+
+    def __post_init__(self) -> None:
+        if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
+            raise ValueError("chunk_index must be a non-negative integer")
+        if not isinstance(self.artifact, ArtifactReference):
+            raise ValueError("artifact must be an ArtifactReference")
+        if not isinstance(self.metrics, Mapping):
+            raise ValueError("metrics must be an object")
+        metrics: dict[str, int | float] = {}
+        for name, value in self.metrics.items():
+            if not isinstance(name, str) or not name:
+                raise ValueError("metric names must be non-empty strings")
+            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
+                raise ValueError("metric values must be finite JSON numbers")
+            metrics[name] = value
+        object.__setattr__(self, "metrics", metrics)
+
+
+@dataclass(frozen=True)
+class FinalResult:
+    """A reducer's durable output, ready for coordinator persistence."""
+
+    artifact: ArtifactReference
+    metrics: Mapping[str, int | float]
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.artifact, ArtifactReference):
+            raise ValueError("artifact must be an ArtifactReference")
+        # Reuse the CompletedPartial metric validation without inventing a fake
+        # artifact lifecycle or widening the result contract.
+        object.__setattr__(self, "metrics", CompletedPartial(0, self.artifact, self.metrics).metrics)
+
+
+def _require_exact_keys(value: Mapping[str, object], expected: set[str], label: str) -> None:
+    actual = set(value)
+    if actual != expected:
+        missing = sorted(expected - actual)
+        unknown = sorted(actual - expected)
+        details: list[str] = []
+        if missing:
+            details.append(f"missing {', '.join(missing)}")
+        if unknown:
+            details.append(f"unknown {', '.join(unknown)}")
+        raise ValueError(f"{label} has {'; '.join(details)} fields")

scimesh/distributed/registry.py

@@ -0,0 +1,96 @@
+"""Registry and orchestration helpers for distributed workload contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Mapping, Sequence
+
+from .models import CompletedPartial, DistributedPlan, FinalResult, _workload_name
+from .workload import DistributedWorkload
+
+
+@dataclass(frozen=True)
+class WorkloadDescription:
+    """Safe metadata that a future coordinator or UI may display."""
+
+    name: str
+    description: str
+
+
+class DistributedWorkloadRegistry:
+    """Collect distributed workloads without coupling them to the CLI registry."""
+
+    def __init__(self) -> None:
+        self._workloads: dict[str, DistributedWorkload] = {}
+
+    def register(self, workload: DistributedWorkload) -> None:
+        name = _workload_name(workload.name)
+        if name in self._workloads:
+            raise ValueError(f"distributed workload already registered: {name}")
+        if not isinstance(workload.description, str) or not workload.description.strip():
+            raise ValueError("distributed workload description must be non-empty")
+        self._workloads[name] = workload
+
+    def require(self, name: str) -> DistributedWorkload:
+        try:
+            return self._workloads[_workload_name(name)]
+        except KeyError as error:
+            raise ValueError(f"unknown distributed workload: {name}") from error
+
+    def descriptions(self) -> tuple[WorkloadDescription, ...]:
+        return tuple(
+            WorkloadDescription(name, workload.description)
+            for name, workload in sorted(self._workloads.items())
+        )
+
+
+class PlanningService:
+    """Small bridge-safe orchestration around a distributed workload registry.
+
+    It writes neither jobs nor artifacts.  A Go coordinator bridge can therefore
+    validate and produce a plan before opening its own all-or-nothing persistence
+    transaction; CTX-08/09 will implement that concrete bridge and reducers.
+    """
+
+    def __init__(self, registry: DistributedWorkloadRegistry) -> None:
+        self._registry = registry
+
+    def plan(
+        self,
+        workload_name: str,
+        input_path: Path,
+        input_artifact_id: str,
+        parameters: Mapping[str, object],
+        shard_rows: int,
+        workspace: Path,
+    ) -> DistributedPlan:
+        if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
+            raise ValueError("shard_rows must be a positive integer")
+        workload = self._registry.require(workload_name)
+        workload.validate_job(parameters)
+        plan = workload.plan(input_path, input_artifact_id, parameters, shard_rows, workspace)
+        if not isinstance(plan, DistributedPlan):
+            raise ValueError("distributed planner must return a DistributedPlan")
+        if plan.workload != workload.name:
+            raise ValueError("distributed planner returned a plan for another workload")
+        # Round-trip through the strict wire schema now, before a future bridge
+        # persists anything. This catches non-JSON values and undeclared fields.
+        return DistributedPlan.from_json(plan.to_json())
+
+    def reduce(
+        self,
+        workload_name: str,
+        partial_results: Sequence[CompletedPartial],
+        parameters: Mapping[str, object],
+        workspace: Path,
+    ) -> FinalResult:
+        workload = self._registry.require(workload_name)
+        indexes = [partial.chunk_index for partial in partial_results]
+        if len(indexes) != len(set(indexes)):
+            raise ValueError("partial results must have unique chunk_index values")
+        ordered = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
+        result = workload.reduce(ordered, parameters, workspace)
+        if not isinstance(result, FinalResult):
+            raise ValueError("distributed reducer must return a FinalResult")
+        return result

scimesh/distributed/workload.py

@@ -0,0 +1,40 @@
+"""Protocol implemented by coordinator-independent distributed workloads."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Mapping, Protocol, Sequence
+
+from .models import CompletedPartial, DistributedPlan, FinalResult
+
+
+class DistributedWorkload(Protocol):
+    """Validate, plan, and reduce one explicit scientific workload.
+
+    ``input_path`` and ``workspace`` are bridge-provided temporary local paths.
+    They must never be included in returned plans or persisted task payloads.
+    """
+
+    name: str
+    description: str
+
+    def validate_job(self, parameters: Mapping[str, object]) -> None:
+        """Reject invalid public parameters before the bridge writes metadata."""
+
+    def plan(
+        self,
+        input_path: Path,
+        input_artifact_id: str,
+        parameters: Mapping[str, object],
+        shard_rows: int,
+        workspace: Path,
+    ) -> DistributedPlan:
+        """Build a JSON-safe plan containing only coordinator artifact references."""
+
+    def reduce(
+        self,
+        partial_results: Sequence[CompletedPartial],
+        parameters: Mapping[str, object],
+        workspace: Path,
+    ) -> FinalResult:
+        """Reduce coordinator-owned partial artifacts in ascending chunk order."""