263 lines
10 KiB
Python
263 lines
10 KiB
Python
"""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")
|