Files
SciMesh/scimesh/distributed/models.py
Emil 6ef92908a1
python / test (push) Waiting to run
Add distributed workload protocol
2026-07-24 14:25:38 +03:00

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")