Merge distributed workload protocol
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
# SciMesh Status
|
||||
|
||||
**Updated:** 2026-07-23
|
||||
**Branch baseline:** `main` at `b4a89dd` (coordinator merge)
|
||||
**Updated:** 2026-07-24
|
||||
**Branch baseline:** `main` at `f953112` (distributed pipeline hardening)
|
||||
|
||||
## Current state
|
||||
|
||||
@@ -32,7 +32,7 @@ Docker PostgreSQL stack on 2026-07-23.
|
||||
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
|
||||
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
|
||||
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
|
||||
| CTX-07 Distributed workload protocol | Not started | Depends on artifact and Worker contracts. |
|
||||
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. The concrete molecular planner/reducer remains CTX-08/09. |
|
||||
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
|
||||
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
|
||||
@@ -41,13 +41,15 @@ Docker PostgreSQL stack on 2026-07-23.
|
||||
|
||||
## Next recommended assignment
|
||||
|
||||
Assign **CTX-07** to the workload role: define distributed job planning and
|
||||
reduction boundaries before implementing distributed search or graph execution.
|
||||
Assign **CTX-08** to the workload role: implement the molecular
|
||||
`similarity-search` planner and worker adapter on top of the accepted CTX-07
|
||||
contract.
|
||||
|
||||
## Known constraints
|
||||
|
||||
- Planner/reducer semantics are not implemented; the operator UI labels
|
||||
`partial_result` files as diagnostic and cannot present them as final output.
|
||||
- The CTX-07 protocol is implemented, but no concrete molecular planner or
|
||||
reducer is registered yet; the operator UI labels `partial_result` files as
|
||||
diagnostic and cannot present them as final output.
|
||||
Use the local `scimesh` CLI for complete workload results.
|
||||
- The worker/coordinator flow currently accepts both underscore API workload
|
||||
names and hyphenated CLI names while the contract is consolidated.
|
||||
|
||||
@@ -2,8 +2,10 @@
|
||||
|
||||
## Status and scope
|
||||
|
||||
This document is the implementation contract for CTX-07. It does not implement
|
||||
a planner, reducer, API endpoint, database migration, or final artifact. Until
|
||||
This document is the implementation contract for CTX-07. Its generic protocol,
|
||||
registry, strict JSON models, and deterministic reduction ordering are
|
||||
implemented in `scimesh/distributed/`. It does not implement a molecular
|
||||
planner, reducer, API endpoint, database migration, or final artifact. Until
|
||||
CTX-08 and CTX-09 are complete, shard CSVs remain diagnostic partial results.
|
||||
|
||||
The protocol gives local scientific workloads a coordinator-independent way to
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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."""
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Contract tests for the coordinator-independent distributed workload boundary."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Mapping, Sequence
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.distributed import (
|
||||
ArtifactReference,
|
||||
CompletedPartial,
|
||||
DistributedPlan,
|
||||
DistributedWorkloadRegistry,
|
||||
FinalResult,
|
||||
PlannedTask,
|
||||
PlanningService,
|
||||
)
|
||||
|
||||
|
||||
def artifact(seed: str, content_type: str = "text/tab-separated-values") -> ArtifactReference:
|
||||
return ArtifactReference(
|
||||
artifact_id=str(uuid5(NAMESPACE_URL, seed)),
|
||||
sha256=(seed.encode("utf-8").hex() * 64)[:64],
|
||||
content_type=content_type,
|
||||
)
|
||||
|
||||
|
||||
class DummyWorkload:
|
||||
"""A deterministic fake workload used to test the generic CTX-07 bridge."""
|
||||
|
||||
name = "dummy-workload"
|
||||
description = "A deterministic test workload."
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.plan_calls = 0
|
||||
self.received_partials: tuple[CompletedPartial, ...] = ()
|
||||
|
||||
def validate_job(self, parameters: Mapping[str, object]) -> None:
|
||||
if parameters != {"mode": "valid"}:
|
||||
raise ValueError("mode must be valid")
|
||||
|
||||
def plan(
|
||||
self,
|
||||
input_path: Path,
|
||||
input_artifact_id: str,
|
||||
parameters: Mapping[str, object],
|
||||
shard_rows: int,
|
||||
workspace: Path,
|
||||
) -> DistributedPlan:
|
||||
self.plan_calls += 1
|
||||
assert input_path.name == "input.tsv"
|
||||
assert workspace.name == "workspace"
|
||||
return DistributedPlan(
|
||||
workload=self.name,
|
||||
resolved_parameters={"mode": parameters["mode"], "source": input_artifact_id},
|
||||
tasks=(
|
||||
PlannedTask(0, artifact(f"{input_artifact_id}:0"), {"mode": "valid"}),
|
||||
PlannedTask(1, artifact(f"{input_artifact_id}:1"), {"mode": "valid"}),
|
||||
),
|
||||
)
|
||||
|
||||
def reduce(
|
||||
self,
|
||||
partial_results: Sequence[CompletedPartial],
|
||||
parameters: Mapping[str, object],
|
||||
workspace: Path,
|
||||
) -> FinalResult:
|
||||
self.received_partials = tuple(partial_results)
|
||||
return FinalResult(artifact("final", "text/csv"), {"partial_count": len(partial_results)})
|
||||
|
||||
|
||||
def service() -> tuple[PlanningService, DummyWorkload]:
|
||||
workload = DummyWorkload()
|
||||
registry = DistributedWorkloadRegistry()
|
||||
registry.register(workload)
|
||||
return PlanningService(registry), workload
|
||||
|
||||
|
||||
def test_unknown_workload_is_rejected_before_a_plan_is_written(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
|
||||
with pytest.raises(ValueError, match="unknown distributed workload"):
|
||||
planner.plan(
|
||||
"unknown-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert workload.plan_calls == 0
|
||||
|
||||
|
||||
def test_invalid_job_is_rejected_before_the_planner_runs(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
|
||||
with pytest.raises(ValueError, match="mode must be valid"):
|
||||
planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
|
||||
{"mode": "invalid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert workload.plan_calls == 0
|
||||
|
||||
|
||||
def test_two_shard_plan_is_deterministic_and_json_serializable(tmp_path: Path) -> None:
|
||||
planner, _ = service()
|
||||
input_artifact_id = artifact("input").artifact_id
|
||||
first = planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
second = planner.plan(
|
||||
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
|
||||
{"mode": "valid"}, 10, tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert first.to_json() == second.to_json()
|
||||
payload = json.loads(first.to_json())
|
||||
assert [task["chunk_index"] for task in payload["tasks"]] == [0, 1]
|
||||
assert all(set(task) == {"chunk_index", "input_artifact", "parameters"} for task in payload["tasks"])
|
||||
assert DistributedPlan.from_json(first.to_json()) == first
|
||||
|
||||
|
||||
def test_plan_rejects_unsafe_or_non_deterministic_task_payloads() -> None:
|
||||
with pytest.raises(ValueError, match="unique, ascending"):
|
||||
DistributedPlan(
|
||||
workload="dummy-workload",
|
||||
resolved_parameters={},
|
||||
tasks=(
|
||||
PlannedTask(1, artifact("one"), {}),
|
||||
PlannedTask(0, artifact("zero"), {}),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="JSON-compatible"):
|
||||
PlannedTask(0, artifact("bad"), {"path": Path("not-serializable")})
|
||||
|
||||
with pytest.raises(ValueError, match="URI or local path"):
|
||||
PlannedTask(0, artifact("uri"), {"input": "file:///tmp/input.tsv"})
|
||||
|
||||
with pytest.raises(ValueError, match="canonical hyphenated"):
|
||||
DistributedPlan("dummy_workload", {}, (PlannedTask(0, artifact("one"), {}),))
|
||||
|
||||
|
||||
def test_reducer_receives_completed_partials_in_chunk_order(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
result = planner.reduce(
|
||||
"dummy-workload",
|
||||
(
|
||||
CompletedPartial(3, artifact("three", "text/csv"), {"scanned_rows": 10}),
|
||||
CompletedPartial(1, artifact("one", "text/csv"), {"scanned_rows": 10}),
|
||||
),
|
||||
{"mode": "valid"},
|
||||
tmp_path / "workspace",
|
||||
)
|
||||
|
||||
assert [partial.chunk_index for partial in workload.received_partials] == [1, 3]
|
||||
assert result.metrics == {"partial_count": 2}
|
||||
|
||||
|
||||
def test_reducer_rejects_duplicate_chunk_indexes_before_invocation(tmp_path: Path) -> None:
|
||||
planner, workload = service()
|
||||
duplicate = CompletedPartial(0, artifact("partial", "text/csv"), {"scanned_rows": 1})
|
||||
|
||||
with pytest.raises(ValueError, match="unique chunk_index"):
|
||||
planner.reduce("dummy-workload", (duplicate, duplicate), {"mode": "valid"}, tmp_path)
|
||||
|
||||
assert workload.received_partials == ()
|
||||
|
||||
|
||||
def test_artifact_references_never_accept_paths_or_uris() -> None:
|
||||
with pytest.raises(ValueError, match="UUID"):
|
||||
ArtifactReference("file:///tmp/input.tsv", "a" * 64, "text/csv")
|
||||
with pytest.raises(ValueError, match="lowercase SHA-256"):
|
||||
ArtifactReference(str(uuid5(NAMESPACE_URL, "input")), "A" * 64, "text/csv")
|
||||
|
||||
|
||||
def test_registry_descriptions_are_stable_and_duplicate_names_are_rejected() -> None:
|
||||
registry = DistributedWorkloadRegistry()
|
||||
first, second = DummyWorkload(), DummyWorkload()
|
||||
registry.register(first)
|
||||
|
||||
assert registry.descriptions()[0].name == "dummy-workload"
|
||||
with pytest.raises(ValueError, match="already registered"):
|
||||
registry.register(second)
|
||||
Reference in New Issue
Block a user