Add distributed workload protocol

2026-07-24 14:25:38 +03:00
parent f953112cfd
commit 6ef92908a1
7 changed files with 625 additions and 9 deletions
@@ -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")